"""长距离精确回忆（联想检索）实验：专门测 ZeroPoint 的已知弱点。

序列格式（每个数字是一个词元）：
    k1 v1 k2 v2 ... kN vN  [问]  kq   →  模型要说出 kq 对应的 v
键从 16 个里不重复地抽 N 个，值从另外 16 个里随机抽。要答对，模型必须把 N 对"键 → 值"原样记住。
注意力天生擅长这个（用 kq 当查询去找那个键）；ZeroPoint 只有 M 个槽，对数一多就会混。

    python3 experiments/recall.py
训练时 N 在 2~16 之间随机，测试时分别报告 N = 2, 4, 8, 12, 16 的准确率。
结果：experiments/recall.md（最新）+ experiments/results/<时间>-recall.md（存档）
"""
import os

import numpy as np

from common import CFG, save_result
from models import TransformerLM, ZeroX
import zerocg as zc
from zerocg import functional as F

ITERS = int(os.environ.get("RECALL_ITERS", 3000))
NK, NV = 16, 16
QUERY = NK + NV          # "问" 这个词元
VOCAB = NK + NV + 1
MAX_N = 16
MAX_LEN = 2 * MAX_N + 2
TEST_N = [2, 4, 8, 12, 16]
SEEDS = [0, 1]
MODELS = {
    "Transformer": lambda: TransformerLM(VOCAB, MAX_LEN, **CFG),
    "ZeroPoint (16 槽)": lambda: ZeroX(VOCAB, n_slot=16, **CFG),
    "ZeroPoint (32 槽)": lambda: ZeroX(VOCAB, n_slot=32, **CFG),
}


def make_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 train_and_eval(make, seed):
    np.random.seed(seed)
    rng = np.random.default_rng(seed)
    model = make()
    opt = zc.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    for it in range(ITERS):
        opt.lr = 3e-3 * min(1, (it + 1) / 100)
        x, y = make_batch(rng, 32, int(rng.integers(2, MAX_N + 1)))
        loss = F.cross_entropy(model(x)[:, -1], y)
        opt.zero_grad()
        loss.backward()
        zc.optim.clip_grad_norm(model.parameters(), 1.0)
        opt.step()
    model.eval()
    test = np.random.default_rng(1000 + seed)
    acc = {}
    with zc.no_grad():
        for n in TEST_N:
            x, y = make_batch(test, 512, n)
            acc[n] = float((model(x).data[:, -1].argmax(-1) == y).mean())
    return acc


def main():
    print(f"联想检索：{len(MODELS)} 个模型 × 种子 {SEEDS}，每次 {ITERS} 步", flush=True)
    res = {}
    for name, make in MODELS.items():
        runs = [train_and_eval(make, s) for s in SEEDS]
        res[name] = {n: np.mean([r[n] for r in runs]) for n in TEST_N}
        print(f"  => {name}: " + "  ".join(f"N={n}:{res[name][n]:.0%}" for n in TEST_N), flush=True)
    lines = ["# 长距离精确回忆（联想检索）\n",
             "`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值，随机猜的准确率是 1/16 ≈ 6%。",
             f"三个模型宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头；训练 {ITERS} 步，N 在 2~{MAX_N} 随机；"
             f"每格是 512 道测试题 × 种子 {SEEDS} 的平均准确率。\n",
             "| 模型 | " + " | ".join(f"N = {n}" for n in TEST_N) + " |", "|---|" + "---|" * len(TEST_N)]
    for name, acc in res.items():
        lines.append(f"| {name} | " + " | ".join(f"{acc[n]:.0%}" for n in TEST_N) + " |")
    save_result("recall", lines)


if __name__ == "__main__":
    main()