"""预注册分析：长期槽到底在干什么？（用 diagnose 保存的 9 个模型，不重训）

    python3 experiments/slots.py

背景：diagnose 发现无 ema 的模型只有 6% 的写入落在半衰期 > 16 字的槽上，但消融说明"分层初始化"有用（去掉变差 0.051）。
结果：experiments/slots.md + results/ 存档。
"""
import os

import numpy as np

from common import HERE, batches, prepare_data, save_result
from models import ZeroX

OUT = os.path.join(HERE, "out")
SEEDS = [0, 1, 2]
T_RUN = 2048
BUCKETS = [(0, 4), (4, 16), (16, 64), (64, 1e9)]   # 半衰期分桶（字）

PREREG = """# 预注册分析：长期槽在干什么

> 写于分析之前（2026-11-01）。跑完只追加结果，不改这一节。用 experiments/out/diag_*.npz（T64 / T64-ema / T256 各 3 种子）。

**矛盾**：T64 模型写入只有 6.1% 落在半衰期 > 16 字的槽；但消融里去掉分层初始化变差 0.051。

**三个可区分的假设**
- H9 读取端在用：长期槽虽然写得少，但**读地址**落在长期槽上的占比明显高于写占比（读占比 ≥ 2 × 写占比）。
  → 长期槽像"慢变的背景"，少量写入、持续读取。
- H10 中期槽才是关键：分层初始化的作用主要来自半衰期 4~64 字的槽，不是 > 64 的；读写都集中在 4~64 桶（该桶读 + 写占比合计 > 50%）。
- H11 训练长度决定长期槽的使用：T256 模型在 > 16 字桶上的写占比 ≥ 2 × T64 的。
  来源：diagnose 里 T256 远处 loss 明显更好，推测它学会了用更长的槽。

**判定**：3 个种子平均；各桶占比按所有层、所有头的平均。H9 / H11 的"2 倍"是无依据的整数阈值，仅作记录；
同时报告原始占比，让读者自己看。
"""


def main():
    text, V, train_data, val_data = prepare_data()
    x, _ = next(batches(val_data, 1, T_RUN, 7, batch=4))
    cfg = dict(n_embd=64, n_head=4, n_layer=2, n_slot=16)
    rows = {}
    for cond in ("T64", "T64-ema", "T256"):
        reads, writes = np.zeros(len(BUCKETS)), np.zeros(len(BUCKETS))
        for seed in SEEDS:
            model = ZeroX(V, ema=(cond == "T64-ema"), **cfg)
            model.load(os.path.join(OUT, f"diag_{cond}_s{seed}.npz")).eval()
            states = model.init_state(4)
            r_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
            w_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
            for t in range(T_RUN):
                _, states = model.step(x[:, t], states, t)
                for l, b in enumerate(model.blocks):
                    a, w = b.mix.last_addr
                    r_acc[l] += a.mean(0)
                    w_acc[l] += w.mean(0)
            for l, b in enumerate(model.blocks):
                hl = np.log(0.5) / np.log(b.mix.retention().data)        # (H, M)
                r = r_acc[l] / r_acc[l].sum(-1, keepdims=True)
                w = w_acc[l] / w_acc[l].sum(-1, keepdims=True)
                for k, (lo, hi) in enumerate(BUCKETS):
                    m = (hl >= lo) & (hl < hi)
                    reads[k] += (r * m).sum() / (r.shape[0]) / len(model.blocks)
                    writes[k] += (w * m).sum() / (w.shape[0]) / len(model.blocks)
        rows[cond] = (reads / len(SEEDS), writes / len(SEEDS))

    lines = [PREREG, "\n## 结果（分析之后写）\n",
             f"验证集 4 条 × {T_RUN} 字递推；占比 = 该桶内所有槽的地址权重之和（读 / 写各自归一化），层和头平均，3 个种子平均。\n",
             "| 条件 | 量 | " + " | ".join(f"半衰期 {lo}~{'∞' if hi > 1e8 else int(hi)}" for lo, hi in BUCKETS) + " |",
             "|---|---|" + "---|" * len(BUCKETS)]
    for cond, (r, w) in rows.items():
        lines.append(f"| {cond} | 读 | " + " | ".join(f"{v:.3f}" for v in r) + " |")
        lines.append(f"| {cond} | 写 | " + " | ".join(f"{v:.3f}" for v in w) + " |")
    r64, w64 = rows["T64"]
    r256, w256 = rows["T256"]
    long_r, long_w = r64[2] + r64[3], w64[2] + w64[3]
    h9 = long_r >= 2 * long_w
    h10 = (r64[1] + r64[2] + w64[1] + w64[2]) / 2 > 0.5
    h11 = (w256[2] + w256[3]) >= 2 * long_w
    lines += ["", "## 判定（按预注册的标准）\n",
              f"- H9 读取端在用长期槽：> 16 字桶 读 {long_r:.3f} vs 写 {long_w:.3f} → **{'成立' if h9 else '不成立'}**（2 倍阈值无依据，仅作记录）",
              f"- H10 中期槽是关键：4~64 字桶 读写平均占比 {(r64[1] + r64[2] + w64[1] + w64[2]) / 2:.3f} > 0.5 → **{'成立' if h10 else '不成立'}**",
              f"- H11 训练长度 256 更用长期槽：> 16 字桶写占比 T256 {w256[2] + w256[3]:.3f} vs T64 {long_w:.3f} → **{'成立' if h11 else '不成立'}**（2 倍阈值无依据，仅作记录）"]
    save_result("slots", lines)


if __name__ == "__main__":
    main()