"""ZeroPoint 消融实验，仿《Attention Is All You Need》表 3：每次只改一处，看那个部件到底有没有用。

    python3 experiments/ablation.py

- 每组 3 个种子，报告 均值 ± 标准差
- 参数量对齐：某个改动让参数量和 base 差超过 1% 时（比如改头数、槽数），自动调整 MLP 宽度把参数量拉回来，
  这样变好变差就不是"参数多了/少了"造成的
- 学习率固定 3e-3（所有变体相同）
结果：experiments/ablation.md（最新）+ experiments/results/<时间>-ablation.md（存档）
"""
import numpy as np

from common import CFG, ITERS, prepare_data, save_result, train
from models import ZeroX

SEEDS = [0, 1, 2]
LR = 3e-3
BASE = dict(n_slot=16, **CFG)

# (行号, 说明, 相对 base 的改动)
VARIANTS = [
    ("base", "ZeroPoint：4 头 × 16 槽，写强度门，写入竞争，分层初始化，λ 可学，可学初始状态，无位置编码", {}),
    ("(A)", "1 个头（每头 64 维）", dict(n_head=1)),
    ("(A)", "8 个头（每头 8 维）", dict(n_head=8)),
    ("(B)", "每头 4 个槽", dict(n_slot=4)),
    ("(B)", "每头 32 个槽", dict(n_slot=32)),
    ("(C)", "去掉写强度门", dict(gate=False)),
    ("(D)", "写入不竞争：softmax 换成每槽独立的 sigmoid", dict(write="sigmoid")),
    ("(E)", "不分层：所有槽初始半衰期相同（约 9.8 个字）", dict(init="uniform")),
    ("(F)", "λ 固定在初始值，不学习", dict(learn_lambda=False)),
    ("(G)", "加上论文 §3.5 的正弦位置编码", dict(pos="sin")),
]


def matched(V, change, target):
    """返回 (参数量对齐后的配置, 是否调整过 MLP)。"""
    cfg = {**BASE, **change}
    n = ZeroX(V, **cfg).num_parameters()
    if abs(n - target) / target <= 0.01:
        return cfg, False
    C, L = cfg["n_embd"], cfg["n_layer"]
    hidden = 4 * C + round((target - n) / (L * (2 * C + 1)))   # 每个隐藏单元每层 2C+1 个参数
    return {**cfg, "mlp_hidden": hidden}, True


def main():
    text, V, train_data, val_data = prepare_data()
    target = ZeroX(V, **BASE).num_parameters()
    print(f"消融实验：{len(VARIANTS)} 组 × 种子 {SEEDS}，每次 {ITERS} 步，lr {LR}", flush=True)
    results = []
    for row, desc, change in VARIANTS:
        cfg, adjusted = matched(V, change, target)
        vals = []
        for seed in SEEDS:
            np.random.seed(seed)
            vals.append(train(f"{row} {desc[:10]}", ZeroX(V, **cfg), train_data, val_data, seed, LR))
        n = ZeroX(V, **cfg).num_parameters()
        results.append((row, desc + (f"（MLP 宽度调为 {cfg['mlp_hidden']} 以对齐参数）" if adjusted else ""),
                        n, np.mean(vals), np.std(vals), vals))
        print(f"  => {row} {desc}: {np.mean(vals):.4f} ± {np.std(vals):.4f}", flush=True)

    base, base_std = results[0][3], results[0][4]
    lines = ["# ZeroPoint 消融实验（仿论文表 3）\n",
             f"每行只改一处，其余和 base 相同。{ITERS} 步，lr {LR}，种子 {SEEDS}，其余训练设置同 [`compare.md`](compare.md)。",
             "参数量和 base 差超过 1% 的变体，自动调整 MLP 宽度对齐参数量。",
             "Δ = 该行验证 loss 均值 − base 均值（**正数表示变差**）。\n",
             f"> 判断标准：base 的种子间标准差是 {base_std:.4f}。|Δ| 小于约 2 倍标准差（{2 * base_std:.4f}）的差别**不能算数**。\n",
             "| | 改动 | 参数量 | 验证 loss | Δ | 结论 |", "|---|---|---|---|---|---|"]
    for row, desc, n, mean, std, _ in results:
        if row == "base":
            lines.append(f"| {row} | {desc} | {n:,} | {mean:.4f} ± {std:.4f} | — | — |")
            continue
        d = mean - base
        noise = 2 * max(base_std, std)
        if noise == 0:
            verdict = "只有 1 个种子，无法判断"
        elif abs(d) < noise:
            verdict = "差别不显著"
        else:
            verdict = "变差：这个部件有用" if d > 0 else "变好：这个部件可能多余"
        lines.append(f"| {row} | {desc} | {n:,} | {mean:.4f} ± {std:.4f} | {d:+.4f} | {verdict} |")
    save_result("ablation", lines)


if __name__ == "__main__":
    main()