import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

# 300 points in 3 well separated blobs. k-means with ONE start (n_init=1),
# repeated with 100 seeds, once with random starting centers and once with
# k-means++. Left: the best solution. Middle: a local minimum it can get stuck
# in. Right: the final inertia of each of the 100 runs.
X, _ = make_blobs(n_samples=300, centers=3, cluster_std=1.0, random_state=42)
runs = {init: [KMeans(n_clusters=3, init=init, n_init=1, random_state=s).fit(X) for s in range(100)]
        for init in ['random', 'k-means++']}
inert = {init: np.array([m.inertia_ for m in ms]) for init, ms in runs.items()}
best = min(runs['k-means++'], key=lambda m: m.inertia_)
stuck = next(m for m in runs['random'] if m.inertia_ > 1.01 * best.inertia_)

BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
col = np.array([BLUE, ORANGE, GREEN])
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 3, figsize=(10, 3.3), gridspec_kw={'width_ratios': [1, 1, 1.15]})
for ax, m, title in [(axes[0], best, 'global minimum'), (axes[1], stuck, 'a local minimum')]:
    ax.scatter(X[:, 0], X[:, 1], s=6, c=col[m.labels_], alpha=0.45)
    ax.scatter(m.cluster_centers_[:, 0], m.cluster_centers_[:, 1], marker='X', s=230,
               c=col, edgecolor='gray', lw=1.6)
    ax.set_title(f'{title}: inertia {m.inertia_:,.1f}', fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
axes[1].text(0.97, 0.03, 'two centers share one blob,\none center covers two',
             transform=axes[1].transAxes, fontsize=8, color='gray', ha='right')

ax = axes[2]
rng = np.random.default_rng(0)
for x, (init, c) in enumerate([('random', ORANGE), ('k-means++', BLUE)]):
    v = inert[init]
    n_bad = (v > 1.01 * v.min()).sum()
    ax.scatter(x + rng.uniform(-0.18, 0.18, len(v)), v, s=12, color=c, alpha=0.7)
    ax.text(x, 9000, f'{n_bad} of 100\nstuck', ha='center', fontsize=8.5, color=c)
ax.set_yscale('log')
ax.set_ylim(400, 16000)
ax.set_xticks([0, 1])
ax.set_xticklabels(['random init', 'k-means++'], fontsize=9)
ax.set_xlim(-0.6, 1.6)
ax.set_ylabel('final inertia (log scale)', fontsize=8.5)
ax.tick_params(labelsize=8)
ax.set_title('100 seeds, one start each (n_init=1)', fontsize=9)
fig.suptitle('Lloyd stops at the nearest minimum: a bad start can cost 10× the inertia', fontsize=10)
plt.tight_layout(rect=(0, 0, 1, 0.95))

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()