import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

# The eight points of the worked example, k = 2, and a deliberately bad start:
# both centers inside the left square. Each panel is one "assign" step, drawn
# with the spokes whose squared lengths add up to the inertia, and an arrow
# showing where the "update" step moves each center next.
X = np.array([[1, 1], [2, 1], [1, 2], [2, 2], [6, 5], [7, 5], [6, 6], [7, 6]], float)
C = np.array([[1, 1], [2, 2]], float)

steps = []
for it in range(3):
    d2 = ((X[:, None, :] - C[None, :, :]) ** 2).sum(axis=2)
    lab = d2.argmin(axis=1)
    inertia = d2[np.arange(len(X)), lab].sum()
    C_new = np.array([X[lab == j].mean(axis=0) for j in range(2)])
    steps.append((C.copy(), lab, inertia, C_new))
    C = C_new

BLUE, ORANGE = '#58a6ff', '#f0883e'
col = [BLUE, ORANGE]
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.4), sharey=True)
titles = ['iteration 0: x₄ joins the wrong group',
          'iteration 1: the centers separate, x₄ returns',
          'iteration 2: nothing changes, converged']
for ax, (Cc, lab, inertia, Cn), title in zip(axes, steps, titles):
    for i, p in enumerate(X):
        c = Cc[lab[i]]
        ax.plot([p[0], c[0]], [p[1], c[1]], color=col[lab[i]], lw=0.9, alpha=0.55, zorder=1)
        ax.scatter(*p, s=46, color=col[lab[i]], zorder=3)
        ax.annotate(f'x{"₁₂₃₄₅₆₇₈"[i]}', p, textcoords='offset points',
                    xytext=(-15 if p[0] in (1, 6) else 9, 6), fontsize=8, color='gray')
    for j in range(2):
        ax.scatter(*Cc[j], marker='X', s=150, color=col[j], edgecolor='gray', lw=0.8, zorder=4)
        if not np.allclose(Cc[j], Cn[j]):
            ax.annotate('', xy=Cn[j], xytext=Cc[j],
                        arrowprops=dict(arrowstyle='->', color='gray', lw=1.2, ls='--'), zorder=2)
    if lab[3] == 1:
        ax.scatter(*X[3], s=260, facecolor='none', edgecolor='#ff7b72', lw=1.6, zorder=5)
    ax.set_title(f'{title}\ninertia = {inertia:.2f}', fontsize=9)
    ax.set_xlim(0, 8)
    ax.set_ylim(0, 7.2)
    ax.set_aspect('equal')
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.2)
axes[0].text(4.1, 0.6, '× = center\n--> next update', fontsize=8, color='gray')
fig.suptitle('Lloyd by hand: assign, update, repeat. The inertia falls 134 → 9.42 → 4.00 and stops',
             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()