# Experiments for "PCA and k-Means": every number in the post comes from here (numpy 2.1, seeded).
#   python unsupervised.py [OUT.json]
import sys, json, math, numpy as np
from usarrests import USARRESTS

def pca(X, scale=True):
    Z = X - X.mean(0); Z = Z / X.std(0, ddof=1) if scale else Z
    U, s, Vt = np.linalg.svd(Z, full_matrices=False)
    for k in range(len(Vt)):                                    # sign convention: largest |loading| positive
        if Vt[k][np.argmax(np.abs(Vt[k]))] < 0: Vt[k] *= -1; U[:, k] *= -1
    return Vt, U * s, s ** 2 / (s ** 2).sum()
def kmeans(X, K, init_labels, max_iter=100):
    lab = init_labels.copy(); hist = []
    for it in range(max_iter):
        C = np.array([X[lab == k].mean(0) if np.any(lab == k) else X[np.random.default_rng(it).integers(len(X))] for k in range(K)])
        new = ((X[:, None, :] - C[None]) ** 2).sum(-1).argmin(1); hist.append(wcss(X, lab, K))
        if np.array_equal(new, lab): break
        lab = new
    return lab, C, wcss(X, lab, K), hist
def wcss(X, lab, K): return float(sum(((X[lab == k] - X[lab == k].mean(0)) ** 2).sum() for k in range(K) if np.any(lab == k)))
def W_pairwise(X, lab, K):   # within-cluster variation as defined by pairwise distances: sum_k (1/|C_k|) sum_{i,i' in C_k} ||x_i - x_i'||^2
    return float(sum((((X[lab == k][:, None] - X[lab == k][None]) ** 2).sum() / np.sum(lab == k)) for k in range(K) if np.any(lab == k)))
def silhouette(X, lab):
    D = np.sqrt(((X[:, None] - X[None]) ** 2).sum(-1)); s = []
    for i in range(len(X)):
        same = lab == lab[i]; a = D[i, same].sum() / max(same.sum() - 1, 1); b = min(D[i, lab == k].mean() for k in set(lab) if k != lab[i]); s.append((b - a) / max(a, b))
    return float(np.mean(s))
def acc(lab, y, K):          # best match of cluster labels to classes (K! permutations)
    import itertools; return max(np.mean(np.array([p[l] for l in lab]) == y) for p in itertools.permutations(range(K)))

rng = np.random.default_rng(0); out = {}
names = [r[0] for r in USARRESTS]; X = np.array([r[1:] for r in USARRESTS], float)
V, S, pve = pca(X); Vr, Sr, pver = pca(X, scale=False)
print("1. USArrests, standardised: PVE", pve.round(4), "cumulative", np.cumsum(pve).round(4))
for k in range(2): print(f"   PC{k + 1} loadings (Murder, Assault, UrbanPop, Rape):", V[k].round(4))
for st in ("Florida", "Nevada", "North Dakota", "Vermont", "Hawaii"): i = names.index(st); print(f"   {st:13s} scores PC1 {S[i, 0]:+.3f} PC2 {S[i, 1]:+.3f}")
print("   unstandardised: PVE", pver.round(4), "PC1 loadings", Vr[0].round(4))
out["usa"] = {"names": names, "X": X.tolist(), "load": V[:2].tolist(), "scores": S[:, :2].round(4).tolist(), "pve": pve.tolist(), "load_raw": Vr[:2].tolist(), "scores_raw": Sr[:, :2].round(4).tolist(), "pve_raw": pver.tolist()}

# 2. the first principal component maximises projected variance and minimises squared distance to the line
Z = X[:, [0, 3]] - X[:, [0, 3]].mean(0); ang = np.linspace(0, np.pi, 3601)          # Murder and Rape: same units (arrests per 100,000), centred
var = [np.var(Z @ np.array([np.cos(a), np.sin(a)]), ddof=1) for a in ang]; dist = [np.sum((Z - np.outer(Z @ np.array([np.cos(a), np.sin(a)]), [np.cos(a), np.sin(a)])) ** 2) for a in ang]
v1 = pca(X[:, [0, 3]], scale=False)[0][0]
print(f"2. Murder and Rape (centred, unscaled): max projected variance at {np.degrees(ang[int(np.argmax(var))]):.2f} deg, min squared distance at {np.degrees(ang[int(np.argmin(dist))]):.2f} deg, PC1 at {np.degrees(np.arctan2(v1[1], v1[0])) % 180:.2f} deg")

out["usa"]["pc1MR"] = float(np.degrees(np.arctan2(v1[1], v1[0])) % 180)
lab0 = np.arange(150) % 3   # a fixed start, used to check the lab's implementation

# 3. k-means: three simulated groups, 50 random starts
C3 = np.array([[0, 0], [3, 0.5], [1.2, 3]]); y3 = np.repeat([0, 1, 2], 50); X3 = C3[y3] + rng.normal(0, 0.8, (150, 2))
res = []
for s in range(50):
    lab, C, w, hist = kmeans(X3, 3, np.random.default_rng(100 + s).integers(0, 3, 150)); res.append((w, len(hist), acc(lab, y3, 3), all(np.diff(hist) <= 1e-9)))
R = np.array([(w, it, a) for w, it, a, _ in res]); best = R[:, 0].min()
print(f"3. k-means K = 3, 50 random starts: best WCSS {best:.3f}, worst {R[:, 0].max():.3f}; starts reaching the best {np.mean(R[:, 0] < best + 1e-6) * 100:.0f}%; "
      f"iterations {int(R[:, 1].min())}-{int(R[:, 1].max())}; objective never increased: {all(m for *_, m in res)}; accuracy of best {R[np.argmin(R[:, 0]), 2]:.3f}")
lab, C, w, _ = kmeans(X3, 3, np.random.default_rng(100).integers(0, 3, 150))
print(f"   pairwise definition W = sum_k (1/|C_k|) sum_(i,i') ||x_i - x_i'||^2 = {W_pairwise(X3, lab, 3):.3f}; 2 x WCSS = {2 * w:.3f}")
out["km"] = {"X": X3.round(4).tolist(), "y": y3.tolist()}
lc, _, wc, _ = kmeans(X3.round(4), 3, lab0); print(f"   check: from labels i mod 3 (rounded data) WCSS {wc:.6f}, cluster sizes {np.bincount(lc).tolist()}")
# hard case: 6 groups, more starts needed
C6 = rng.uniform(0, 10, (6, 2)); y6 = np.repeat(np.arange(6), 30); X6 = C6[y6] + rng.normal(0, 0.7, (180, 2))
w6 = [kmeans(X6, 6, np.random.default_rng(500 + s).integers(0, 6, 180))[2] for s in range(100)]
print(f"   K = 6 (six groups), 100 random starts: best {min(w6):.2f}, worst {max(w6):.2f}, starts within 0.1% of best {np.mean(np.array(w6) < min(w6) * 1.001) * 100:.0f}%")

# 4. choosing K: WCSS always falls; silhouette
for K in range(2, 7):
    ws = [kmeans(X3, K, np.random.default_rng(900 + s).integers(0, K, 150)) for s in range(20)]; lb, _, wk, _ = min(ws, key=lambda r: r[2])
    print(f"4. K = {K}: best WCSS {wk:.2f}, silhouette {silhouette(X3, lb):.3f}")

# 5. how many partitions: Stirling numbers of the second kind
def stirling2(n, k): return sum((-1) ** j * math.comb(k, j) * (k - j) ** n for j in range(k + 1)) // math.factorial(k)
for n, k in ((10, 3), (20, 4), (150, 3)): print(f"5. partitions of {n} points into {k} clusters: {stirling2(n, k):.4e} (K^n = {k ** n:.4e}, K^n/K! = {k ** n / math.factorial(k):.4e})")

# 6. correlation 1 vs 0 and the centroids
v = rng.normal(0, 1, 300); Xa = np.c_[v, 2 * v + 1]; la = kmeans(Xa, 3, rng.integers(0, 3, 300))[1]
Xb = rng.normal(0, 1, (300, 2)); lb_ = kmeans(Xb, 3, rng.integers(0, 3, 300))[1]
col = lambda C: abs(np.linalg.det(np.c_[C[1] - C[0], C[2] - C[0]]))
print(f"6. corr(V1, V2) = 1: centroids collinear, |det| = {col(la):.2e}; corr = {np.corrcoef(Xb.T)[0, 1]:.3f}: |det| = {col(lb_):.3f}")

# 7. where k-means struggles: scale, shape, unequal spread
y = np.repeat([0, 1], 100); Xs = np.c_[rng.normal(0, 1, 200) + 3 * y, rng.normal(0, 1000, 200)]
a_raw = acc(min((kmeans(Xs, 2, np.random.default_rng(s).integers(0, 2, 200)) for s in range(20)), key=lambda r: r[2])[0], y, 2)
Zs = (Xs - Xs.mean(0)) / Xs.std(0); a_std = acc(min((kmeans(Zs, 2, np.random.default_rng(s).integers(0, 2, 200)) for s in range(20)), key=lambda r: r[2])[0], y, 2)
print(f"7. income-like second feature in large units: accuracy raw {a_raw:.3f}, standardised {a_std:.3f}")
t = rng.uniform(0, np.pi, 100); Xm = np.r_[np.c_[np.cos(t), np.sin(t)], np.c_[1 - np.cos(t), 0.5 - np.sin(t)]] + rng.normal(0, 0.08, (200, 2))
a_m = acc(min((kmeans(Xm, 2, np.random.default_rng(s).integers(0, 2, 200)) for s in range(20)), key=lambda r: r[2])[0], y, 2)
yu = np.r_[np.zeros(170, int), np.ones(30, int)]; Xu = np.r_[rng.normal(0, 2.0, (170, 2)), rng.normal(0, 0.3, (30, 2)) + [4.5, 0]]
a_u = acc(min((kmeans(Xu, 2, np.random.default_rng(s).integers(0, 2, 200)) for s in range(20)), key=lambda r: r[2])[0], yu, 2)
print(f"   two moons: accuracy {a_m:.3f}; one wide and one small tight cluster: accuracy {a_u:.3f}")
out["fail"] = {"moons": Xm.round(4).tolist(), "uneq": Xu.round(4).tolist(), "y_uneq": yu.tolist()}

# 8. cluster IDs as a feature: ordinal codes vs dummies
Xc = np.r_[rng.normal([0, 0], 0.5, (60, 2)), rng.normal([3, 0], 0.5, (60, 2)), rng.normal([1.5, 3], 0.5, (60, 2))]; g = np.repeat([0, 1, 2], 60)
yv = np.array([2.0, -1.0, 4.0])[g] + rng.normal(0, 1, 180); labc = kmeans(Xc, 3, rng.integers(0, 3, 180))[0]
import itertools
r2 = lambda Z: 1 - np.sum((yv - Z @ np.linalg.lstsq(Z, yv, rcond=None)[0]) ** 2) / np.sum((yv - yv.mean()) ** 2)
ords = [r2(np.c_[np.ones(180), np.array(p)[labc]]) for p in itertools.permutations(range(3))]
dum = r2(np.c_[np.ones(180), labc == 1, labc == 2])
print(f"8. regression on the cluster label: as one ordinal number, R^2 ranges {min(ords):.3f}-{max(ords):.3f} over the 6 possible numberings; as dummies {dum:.3f} for any numbering")
if len(sys.argv) > 1: json.dump(out, open(sys.argv[1], "w"), separators=(",", ":"))