"""精确回忆实验的诊断脚本（诊断性，不参与预注册判定）。

回答的问题：为什么 recall.py 跑 3000 步和 20000 步结果几乎一样（Transformer N=16 只有 14~18%）？

    python3 experiments/recall_diag.py

四段，都对固定测试题评估、种子 0：
  A 固定 32 条 N=2 题反复训练：能不能背下来（检验框架/优化）。
  B 只训 N=1（复制值，不需要匹配）：能不能学会复制。
  C 只训 N=2，比较学习率安排：恒定 3e-3 / 余弦 3e-3→3e-4 / 余弦 1e-3→1e-4，各 8000 步。
  D 混合 N（2~16）用 1e-3 余弦跑 12000 步：看混合训练能不能学会。

结论（11-08，见 recall.md 的诊断一节）：A、B 都能 100%；C 里只有 1e-3 余弦能学会 N=2
（第 8000 步 100%，loss 0.001），恒定 3e-3 跑 20000 步也停在 53%（loss ≈ ln2）；D 12000 步仍 56%，
因为混合数据里 N=2 只占约 1/15，信号被稀释。
"""
import sys
import os
import time

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

import zerocg as zc  # noqa: E402
from zerocg import functional as F  # noqa: E402
from models import TransformerLM  # noqa: E402

NK, NV = 16, 16
QUERY = NK + NV
CFG = dict(n_embd=64, n_head=4, n_layer=2)
MAX_N = 16
TEST_N = [2, 4, 8, 12, 16]


def batch(rng, B, n):
    keys = np.stack([rng.permutation(NK)[:n] for _ in range(B)])
    vals = rng.integers(NK, NK + NV, size=(B, n))
    seq = np.empty((B, 2 * n + 2), dtype=np.int64)
    seq[:, 0:2 * n:2], seq[:, 1:2 * n:2] = keys, vals
    q = rng.integers(0, n, B)
    seq[:, -2], seq[:, -1] = QUERY, keys[np.arange(B), q]
    return seq, vals[np.arange(B), q]


def model():
    return TransformerLM(NK + NV + 1, 2 * MAX_N + 2, **CFG)


def acc(model, tests):
    with zc.no_grad():
        return {n: round(float((model(tests[n][0]).data[:, -1].argmax(-1) == tests[n][1]).mean()), 3)
                for n in tests}


def tests():
    rng = np.random.default_rng(7000)
    return {n: batch(rng, 512, n) for n in TEST_N}


def run(label, iters, lr0, lr1, warmup, n_fixed=None, seed=0, eval_every=2000):
    np.random.seed(seed)
    rng = np.random.default_rng(seed)
    m = model()
    opt = zc.optim.AdamW(m.parameters(), lr=lr0, weight_decay=0.01)
    t = tests()
    t0 = time.time()
    for it in range(iters):
        frac = it / max(1, iters - 1)
        opt.lr = lr0 * min(1, (it + 1) / max(1, warmup)) * (lr1 / lr0) ** frac
        n = n_fixed if n_fixed else int(rng.integers(2, MAX_N + 1))
        x, y = batch(rng, 32, n)
        loss = F.cross_entropy(m(x)[:, -1], y)
        opt.zero_grad()
        loss.backward()
        zc.optim.clip_grad_norm(m.parameters(), 1.0)
        opt.step()
        if (it + 1) % eval_every == 0:
            a = acc(m, t)
            print(f"  {label} step {it+1:5d} loss {float(loss.data):.3f} "
                  + "  ".join(f"N={k}:{a[k]}" for k in TEST_N) + f"  ({time.time()-t0:.0f}s)", flush=True)
    return a


def main():
    print("A 固定 32 条 N=2 反复训练（记忆）", flush=True)
    np.random.seed(0)
    m = model()
    rng = np.random.default_rng(0)
    x, y = batch(rng, 32, 2)
    opt = zc.optim.AdamW(m.parameters(), lr=1e-3, weight_decay=0.01)
    for it in range(600):
        opt.lr = 3e-3 * min(1, (it + 1) / 100)
        loss = F.cross_entropy(m(x)[:, -1], y)
        opt.zero_grad()
        loss.backward()
        opt.step()
        if (it + 1) % 200 == 0:
            with zc.no_grad():
                a = float((m(x).data[:, -1].argmax(-1) == y).mean())
            print(f"  mem step {it+1} loss {float(loss.data):.4f} acc {a}", flush=True)

    print("B 只训 N=1（复制）", flush=True)
    run("N1", 1500, 3e-3, 3e-3, 100, n_fixed=1, eval_every=500)

    print("C1 只训 N=2，恒定 3e-3（recall.py 的原设置）", flush=True)
    run("C1", 8000, 3e-3, 3e-3, 100, n_fixed=2)
    print("C2 只训 N=2，余弦 3e-3→3e-4", flush=True)
    run("C2", 8000, 3e-3, 3e-4, 200, n_fixed=2)
    print("C3 只训 N=2，余弦 1e-3→1e-4", flush=True)
    run("C3", 8000, 1e-3, 1e-4, 200, n_fixed=2)

    print("D 混合 N，余弦 1e-3→1e-4，12000 步", flush=True)
    run("D", 12000, 1e-3, 1e-4, 200, n_fixed=None)


if __name__ == "__main__":
    main()