"""预注册实验：ZeroPoint 长序列变差的原因是不是槽值无界累积？ema 写入能不能修好？

    python3 experiments/ema_longrun.py

实验设计在跑之前写定（见 PREREG），结果出来后不改判定标准。结果：experiments/ema_longrun.md + results/ 存档。
"""
import numpy as np

from common import CFG, ITERS, T, batches, evaluate, prepare_data, save_result, train
from models import ZeroX
import zerocg as zc

SEEDS = [0, 1, 2]
LR = 6e-3                      # compare.md 里 ZeroPoint 的最佳学习率
LONG_T = 4096
LONG_N = 32                    # 外推测试序列条数（compare.py 只有 8 条，这里加到 32）
SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)]
PROBE_POS = [64, 256, 1024, 4000]   # 测槽值范围的位置

PREREG = """# 预注册实验：长序列变差与 ema 写入

> 写于实验之前（2026-10-18）。跑完后只在下面追加结果，不改这一节。

**观察**：compare.md 里 ZeroPoint（无 ema）在 1024~4095 字处的 loss（1.70）高于 64~255 字处（1.62）。

**假设 H1**：原因是长期槽（λ 接近 1）的数值随序列长度无界累积，超出训练时（64 字）见过的范围。

**预测**：
- ① 直接测：推理时槽值的最大绝对值，无 ema 的模型在 4000 字处明显大于 64 字处（比值 > 2）；有 ema 的模型比值接近 1（< 1.5）。
- ② 间接测：开 ema 重新训练（其余完全相同），1024~4095 字处的 loss 不再明显高于 64~255 字处：差值 < 2 × 种子间标准差。

**H1 不成立时会看到**：① 无 ema 的槽值没有明显增长；或 ② 开 ema 后远处 loss 照样升高。

**判定**：① 和 ② 都成立 → 接受 H1，且 ema 是有效的修法；只有 ② 成立 → ema 有效但原因另有其他；
都不成立 → 放弃 H1，下一步查学到的 λ 值。另外要求 ema 不损害短距离：64 字处验证 loss 的差别 < 2 × 标准差。

**对照**：两组用同样的 3 个种子、同样 {iters} 步、学习率 {lr}、同样语料；外推用 {n} 条 {lt} 字的验证序列，
两组都用分块并行算法算 loss（compare.py 里 ZeroPoint 用的是递推算法，和 Transformer 不对等，这里统一）。
"""


def long_loss(model, val_data):
    """(各段 loss 均值) ；分 4 批算，省内存。"""
    nll_all = []
    for x, y in batches(val_data, 4, LONG_T, 7, batch=LONG_N // 4):
        with zc.no_grad():
            logits = model(x).data
        logp = logits - logits.max(-1, keepdims=True)
        logp = logp - np.log(np.exp(logp).sum(-1, keepdims=True))
        nll_all.append(-np.take_along_axis(logp, y[..., None], -1)[..., 0])
    nll = np.concatenate(nll_all)
    return np.array([nll[:, a:b].mean() for a, b in SEGMENTS])


def slot_magnitude(model, val_data):
    """递推跑 4096 字，记录各位置处所有层槽值的最大绝对值。返回 {位置: 最大|S|}。"""
    x, _ = next(batches(val_data, 1, LONG_T, 7, batch=4))
    states = model.init_state(4)
    out = {}
    for t in range(LONG_T):
        _, states = model.step(x[:, t], states, t)
        if t + 1 in PROBE_POS:
            out[t + 1] = max(float(np.abs(S).max()) for S in states)
    return out


def main():
    text, V, train_data, val_data = prepare_data()
    print(PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), flush=True)
    res = {}
    for ema in (False, True):
        name = "ema" if ema else "no-ema"
        vals, longs, mags = [], [], []
        for seed in SEEDS:
            np.random.seed(seed)
            model = ZeroX(V, n_slot=16, ema=ema, **CFG)
            vals.append(train(name, model, train_data, val_data, seed, LR))
            model.eval()
            longs.append(long_loss(model, val_data))
            mags.append(slot_magnitude(model, val_data))
            print(f"  => {name} seed={seed}: val {vals[-1]:.4f}; long {np.round(longs[-1], 4)}; "
                  f"|S| {[round(mags[-1][p], 2) for p in PROBE_POS]}", flush=True)
        res[name] = dict(val=np.array(vals), long=np.array(longs),
                         mag=np.array([[m[p] for p in PROBE_POS] for m in mags]))

    # ---- 结果 + 按预注册的标准判定 ----
    lines = [PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), "\n## 结果（实验之后写）\n",
             "### 64 字处的验证 loss（均值 ± 标准差，3 个种子）\n", "| 模型 | 验证 loss |", "|---|---|"]
    for name, r in res.items():
        lines.append(f"| {name} | {r['val'].mean():.4f} ± {r['val'].std():.4f} |")
    lines += ["", f"### 长度外推（{LONG_N} 条 × {LONG_T} 字，分块并行算法；均值 ± 标准差）\n",
              "| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)]
    for k, (a, b) in enumerate(SEGMENTS):
        lines.append(f"| {a} ~ {b - 1} | " + " | ".join(f"{r['long'][:, k].mean():.4f} ± {r['long'][:, k].std():.4f}"
                                                        for r in res.values()) + " |")
    lines += ["", "### 槽值的最大绝对值（递推推理，3 个种子平均）\n",
              "| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)]
    for k, p in enumerate(PROBE_POS):
        lines.append(f"| {p} | " + " | ".join(f"{r['mag'][:, k].mean():.2f}" for r in res.values()) + " |")

    no, em = res["no-ema"], res["ema"]
    ratio_no = (no["mag"][:, 3] / no["mag"][:, 0]).mean()
    ratio_em = (em["mag"][:, 3] / em["mag"][:, 0]).mean()
    p1 = ratio_no > 2 and ratio_em < 1.5
    far, near = em["long"][:, 3], em["long"][:, 1]
    gap, noise = (far - near).mean(), 2 * max(far.std(), near.std())
    p2 = gap < noise
    gap_no = (no["long"][:, 3] - no["long"][:, 1]).mean()
    short_ok = abs(em["val"].mean() - no["val"].mean()) < 2 * max(em["val"].std(), no["val"].std())
    if p1 and p2:
        verdict = "① ② 都成立：**接受 H1**，ema 是有效的修法"
    elif p2:
        verdict = "只有 ② 成立：ema 有效，但长序列变差的原因不是（或不只是）槽值无界"
    elif p1:
        verdict = "只有 ① 成立：槽值确实无界，但 ema 没有修好远处的 loss，H1 不足以解释变差"
    else:
        verdict = "① ② 都不成立：**放弃 H1**，下一步查学到的 λ 值"
    lines += ["", "## 判定（按预注册的标准）\n",
              f"- 预测 ①（槽值比值 4000 字 / 64 字）：无 ema {ratio_no:.2f}，有 ema {ratio_em:.2f} → **{'成立' if p1 else '不成立'}**",
              f"- 预测 ②（ema 模型 1024~4095 与 64~255 的 loss 差）：{gap:+.4f}，噪声阈值 {noise:.4f} → **{'成立' if p2 else '不成立'}**"
              f"（对照：无 ema 的差是 {gap_no:+.4f}）",
              f"- ema 不损害短距离：64 字处 loss 差 {em['val'].mean() - no['val'].mean():+.4f} → **{'成立' if short_ok else '不成立'}**",
              f"\n**结论**：{verdict}" + ("" if short_ok else "；但 ema 损害了短距离效果，不能直接采用")]
    save_result("ema_longrun", lines)


if __name__ == "__main__":
    main()