"""
=============================================================================
第 5 讲 · 代码 1：k-means —— 坐标下降、局部最优、k-means++、怎样选 k、它假设了什么
=============================================================================

对应讲义：lectures/lecture05_unsupervised.md 第 2 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
k-means 大概是你学过的第一个聚类算法。复习时值得问的是"它到底在优化什么、为什么会卡住、什么时候根本就不该用它"：

    演示 1：k-means = distortion J(c, μ) 上的坐标下降。逐轮打印 J（assignment 之后、update 之后），
            验证它单调不增，并在有限步内停下。
    演示 2：局部最优。同一份数据（3×3 = 9 个分得很开的团），200 次随机初始化 vs 200 次 k-means++：
            最终 J 的分布、多少次找到最好的解、"重启 10 次取最好"的成功率。
    演示 3：怎样选 k：distortion 随 k 单调下降（嵌套的模型族必然如此），所以不能直接选 J 最小的 k；
            看"肘部"（elbow）—— 团分得开时肘部很清楚，团重叠时就模糊了。
    演示 4：k-means 隐含的假设（球形、大小相近的簇）。在细长的簇和大小悬殊的簇上，
            k-means 的解比"真实的划分"有**更低**的 J —— 不是优化失败，是目标函数本身就偏好错误的答案。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_kmeans.py          # 约 2 秒

依赖本目录的 unsupervised.py（它会导入第 1 讲的 datasets.py 和 linear_models.py）。
"""

import numpy as np

from unsupervised import (best_perm_accuracy, distortion, init_kmeanspp, init_random, kmeans, kmeans_best_of,
                          make_grid_blobs, make_spec, print_header)

GRID = dict(n_per=60, side=3, spacing=4.0, std=0.7)     # 演示 1~3 共用的数据：9 个团，每团 60 个点


def grid_data(std=GRID["std"], seed=11):
    """生成 3×3 网格高斯团数据。固定种子 → 演示 1~3 用的是同一份数据（演示 3 还会换一个更大的 std）。"""
    return make_grid_blobs(GRID["n_per"], GRID["side"], GRID["spacing"], std, np.random.default_rng(seed))


# =============================================================================
# 演示 1：坐标下降，J 单调不增
# =============================================================================
def demo_coordinate_descent():
    print_header("演示 1：k-means = coordinate descent on the distortion J(c, μ)")
    X, y = grid_data()
    rng = np.random.default_rng(12)
    hist = []
    C, labels, J, n_iters = kmeans(X, init_random(X, 9, rng), history=hist)
    print(f"数据：{len(X)} 个点，9 个团（3×3 网格，间距 {GRID['spacing']:g}，每团标准差 {GRID['std']:g}）；K = 9，随机初始化\n")
    print(f"{'iter':>5}{'J after assign':>17}{'J after update':>17}{'decrease':>12}")
    prev = None
    for t, (Ja, Ju) in enumerate(hist, start=1):
        dec = "" if prev is None else f"{prev - Ju:.1f}"
        print(f"{t:>5}{Ja:>17.2f}{Ju:>17.2f}{dec:>12}")
        prev = Ju
    seq = [j for pair in hist for j in pair]
    monotone = all(b <= a + 1e-9 * a for a, b in zip(seq, seq[1:]))
    assert monotone, "distortion increased — this should never happen"
    blob_centers = np.array([X[y == k].mean(axis=0) for k in range(9)])
    J_truth = distortion(X, blob_centers, y)
    acc = best_perm_accuracy(y, labels, 9)
    # 每个学到的中心离哪个真实团最近？数一数每个团分到了几个中心（理想情况：每团恰好 1 个）
    nearest_blob = np.argmin(((C[:, None, :] - blob_centers[None]) ** 2).sum(axis=2), axis=1)
    per_blob = np.bincount(nearest_blob, minlength=9)
    print(f"\n  {n_iters} 轮后分配不再变化，停止。J 单调不增：{monotone}（已用 assert 检查每个半步）")
    print(f"  最终 J = {J:.2f}；用真实的团编号划分时 J = {J_truth:.2f}；和真实团的匹配准确率 = {acc:.3f}")
    print(f"  每个真实团分到的中心个数 = {per_blob.tolist()}："
          f"{int(np.sum(per_blob >= 2))} 个团挤了 2 个以上中心，{int(np.sum(per_blob == 0))} 个团没有自己的中心")
    if J > 1.01 * J_truth:
        stuck = (f"这次随机初始化陷入了局部最优（J = {J:.0f}，真实划分 {J_truth:.0f}）：有的团挤着两个中心，"
                 "于是别处必然有一个中心要同时管两个团")
    else:
        stuck = "这次随机初始化找到了真实的结构"
    first, last = hist[0], hist[-1]
    print(f"""
观察 Observations：
  · 每个半步都是在另一组变量固定时**精确地**最小化 J：assignment 让每个点去最近的中心，update 把中心移到簇的均值。
    所以 J 只降不升，而分配方式只有有限多种，算法一定在有限步内停下（这里 {n_iters} 轮）。
  · 第 1 轮就把 J 从 {first[0]:.0f} 降到 {first[1]:.0f}；最后两轮只降了 {hist[-3][1] - hist[-2][1]:.1f} 和 {hist[-2][1] - last[1]:.1f} —— 只是边界上的几个点换来换去。
  · 停下 ≠ 全局最优：{stuck}。
    演示 2 统计这件事发生得有多频繁。""")


# =============================================================================
# 演示 2：局部最优；随机初始化 vs k-means++
# =============================================================================
def demo_local_optima():
    print_header("演示 2：local optima —— 200 random inits vs 200 k-means++ inits")
    X, y = grid_data()
    n_runs = 200
    results = {}
    for name, init_fn, seed in [("random (Forgy)", init_random, 21), ("k-means++", init_kmeanspp, 22)]:
        rng = np.random.default_rng(seed)
        J0s, Js, its = [], [], []
        for _ in range(n_runs):
            C0 = init_fn(X, 9, rng)
            J0s.append(distortion(X, C0, np.argmin(((X[:, None, :] - C0[None]) ** 2).sum(axis=2), axis=1)))
            _, _, J, it = kmeans(X, C0)
            Js.append(J)
            its.append(it)
        results[name] = (np.array(J0s), np.array(Js), np.array(its))
    J_best = min(r[1].min() for r in results.values())
    J_truth = distortion(X, np.array([X[y == k].mean(axis=0) for k in range(9)]), y)
    print(f"9 个团，K = 9。两种初始化各跑 {n_runs} 次。全部 400 次里最好的 J* = {J_best:.2f}"
          f"（真实划分的 J = {J_truth:.2f}）\n")
    print(f"{'init':<16}{'J at init':>11}{'best J':>10}{'median J':>10}{'90% J':>10}{'worst J':>10}"
          f"{'P(J<=1.01J*)':>14}{'iters':>7}{'best-of-10 ok':>15}")
    summary = {}
    for name, (J0s, Js, its) in results.items():
        hit = Js <= 1.01 * J_best
        groups = hit.reshape(20, 10).any(axis=1)          # 把 200 次按顺序分成 20 组，每组 10 次：组内有一次命中就算成功
        summary[name] = (hit.mean(), groups.mean(), np.median(Js))
        print(f"{name:<16}{np.median(J0s):>11.1f}{Js.min():>10.1f}{np.median(Js):>10.1f}"
              f"{np.quantile(Js, 0.9):>10.1f}{Js.max():>10.1f}{hit.mean():>14.3f}{its.mean():>7.1f}"
              f"{groups.mean():>15.2f}")
    print("  （J at init = 初始中心不动、只做一次分配时的 J，取中位数；iters = 平均轮数；"
          "best-of-10 ok = 20 组“重启 10 次取最好”中找到 ≤1.01·J* 的比例）")
    pr, pp = summary["random (Forgy)"], summary["k-means++"]
    print(f"""
观察 Observations：
  · 同一份分得很开的数据，随机初始化只有 {pr[0]:.0%} 的运行找到最好的解；中位数 J 是 {pr[2]:.0f}，约为 J* 的 {pr[2] / J_best:.1f} 倍。
    典型的局部最优：两个中心挤在同一个团里，另一个中心横跨两个团 —— 单个中心挪不出去，因为挪动任何一个都会让 J 变大。
  · k-means++ 把成功率提高到 {pp[0]:.0%}，中位数 J 就是 J*；它的初始 J（还没迭代）就已经比随机初始化低得多，
    而且需要的迭代轮数也更少。但它不是万能的：仍有 {1 - pp[0]:.0%} 的运行卡在局部最优。
  · 所以实践中总是**多次重启取 J 最小的**（sklearn 默认 n_init 次）。重启 10 次时，随机初始化的成功率是 {pr[1]:.0%}，
    k-means++ 是 {pp[1]:.0%}。
  · J 是非凸的（分配变量是离散的），找全局最优是 NP-hard；k-means++ 只保证期望上 O(log K) 倍以内。""")


# =============================================================================
# 演示 3：怎样选 k
# =============================================================================
def demo_choose_k():
    print_header("演示 3：choosing k —— distortion always decreases; look for the elbow")
    datasets = {"well separated (std 0.7)": grid_data(std=0.7), "overlapping (std 1.6)": grid_data(std=1.6)}
    ks = list(range(1, 13))
    table = {}
    for name, (X, _) in datasets.items():
        rng = np.random.default_rng(31)
        table[name] = [kmeans_best_of(X, k, rng, n_init=10)[2] for k in ks]
    names = list(datasets)
    print(f"每个 k 用 k-means++ 重启 10 次取最好。ratio = J(k−1)/J(k)：这一步让 J 缩小了多少倍\n")
    print(f"{'k':>3}" + "".join(f"{'J (' + nm.split()[0] + ')':>18}{'ratio':>8}" for nm in names))
    for i, k in enumerate(ks):
        row = f"{k:>3}"
        for nm in names:
            J = table[nm][i]
            ratio = "" if i == 0 else f"{table[nm][i - 1] / J:.2f}"
            row += f"{J:>18.1f}{ratio:>8}"
        print(row)
    sep, ovl = table[names[0]], table[names[1]]
    r_sep = [sep[i - 1] / sep[i] for i in range(1, len(ks))]
    r_ovl = [ovl[i - 1] / ovl[i] for i in range(1, len(ks))]
    print(f"""
观察 Observations：
  · 两列 J 都随 k 单调下降：k = n 时每个点自成一簇，J = 0。所以"选 J 最小的 k"毫无意义 —— 和第 1 讲
    "多项式次数越高、训练误差越低"是同一件事（嵌套的模型族）。
  · 分得开的数据：k = 8 → 9 时 J 缩小 {r_sep[7]:.2f} 倍，而 k = 9 → 10 只缩小 {r_sep[8]:.2f} 倍 —— 清楚的肘部（elbow），
    正好在真实的团数 9。（前面几步的比值也很大，因为每多一个中心就"解决"一个或几个团。）
  · 重叠的数据：同样两步的比值是 {r_ovl[7]:.2f} 和 {r_ovl[8]:.2f}，曲线平滑地下降，看不出明确的肘部。
    真实数据通常更像这一列。
  · 所以选 k 往往要靠别的依据：下游任务的效果；GMM 的 held-out 似然或 BIC（代码 02 的模型有概率解释）；
    silhouette、gap statistic 等启发式；或者干脆由应用决定（"把客户分成 5 档"）。""")


# =============================================================================
# 演示 4：k-means 隐含的假设
# =============================================================================
def demo_assumptions():
    print_header("演示 4：what k-means assumes —— elongated and unequal-size clusters")
    print(f"{'data':<12}{'J (k-means)':>13}{'J (true split)':>16}{'accuracy':>10}"
          f"{'pts moved':>11}  {'what k-means did'}")
    notes = {}
    for name, seed in [("elongated", 41), ("unequal", 42)]:
        X, y = make_spec(name, np.random.default_rng(seed))
        C, labels, J = kmeans_best_of(X, 2, np.random.default_rng(seed + 100), n_init=10)
        C_true = np.array([X[y == k].mean(axis=0) for k in range(2)])
        J_true = distortion(X, C_true, y)
        acc = best_perm_accuracy(y, labels, 2)
        # 被"挪走"的点数：最好的标签匹配下，标签和真实不一致的点
        moved = int(round((1 - acc) * len(y)))
        if name == "elongated":
            # 两个中心的连线方向：如果主要沿 x 方向，说明 k-means 把两根长条"竖着一刀切成左右两半"
            dx, dy = np.abs(C[0] - C[1])
            what = f"centers differ mainly in {'x' if dx > dy else 'y'} (|dx|={dx:.2f}, |dy|={dy:.2f})"
        else:
            big = labels[y == 0]                              # 大簇的 400 个点被分到了哪里
            small_label = np.bincount(labels[y == 1], minlength=2).argmax()
            what = f"{int(np.sum(big == small_label))} of 400 big-cluster points go to the small cluster"
        notes[name] = (J, J_true, acc)
        print(f"{name:<12}{J:>13.1f}{J_true:>16.1f}{acc:>10.3f}{moved:>11}  {what}")
    e, u = notes["elongated"], notes["unequal"]
    print(f"""
观察 Observations：
  · 细长的簇：k-means 没有把上下两根长条分开，而是把它们一起"左右切开"，准确率只有 {e[2]:.3f}（随机猜是 0.5）。
    关键是 J 的比较：k-means 的解 J = {e[0]:.0f}，真实划分 J = {e[1]:.0f} —— **错误的答案 J 更低**。
    每根长条在 x 方向的方差很大（16），沿 x 切一刀比沿 y 切一刀减少的平方距离多得多。
  · 大小悬殊的簇：k-means 的边界是两个中心的垂直平分线，它不知道大簇"很宽"、小簇"很紧"，
    于是把大簇靠近小簇的一侧划给了小簇。同样，k-means 的 J（{u[0]:.0f}）低于真实划分（{u[1]:.0f}）。
  · 结论：这不是优化失败（多次重启、k-means++ 都救不了），而是 J 本身隐含的假设 ——
    每个簇都是"半径差不多的球"。k-means 正是 Σ_k = σ²I、权重相等的 GMM 在 σ → 0 时的 EM（代码 02 演示 3）。
    放松这些假设（每个簇自己的协方差和权重）就得到 GMM + EM：代码 02 演示 3 在同样的两份数据上比较它和 Bayes 最优。
  · 实践中的补救：先做合适的特征变换 / 标准化；或者换 GMM、谱聚类（spectral clustering）、DBSCAN 等。""")


def main():
    demo_coordinate_descent()
    demo_local_optima()
    demo_choose_k()
    demo_assumptions()


if __name__ == "__main__":
    main()