"""ZeroPoint vs Transformer 对照实验（字符级语言模型）。

    python3 experiments/compare.py            # 默认语料：Python 标准库源码（离线可用）
    python3 experiments/compare.py my.txt     # 换成自己的文本
    ITERS=2000 python3 experiments/compare.py

三个模型宽度、层数、头数、MLP 完全一样，唯一区别是"token 混合层 + 位置编码"：
    Transformer        自注意力 + 可学习位置编码
    Transformer (sin)  自注意力 + 论文 §3.5 的正弦位置编码（理论上能外推到更长序列）
    ZeroPoint            记忆槽，无位置编码

流程：
    1. 学习率搜索：每个模型在 LRS 里各试一次（种子 0），取验证 loss 最低的
    2. 用各自最好的学习率再跑种子 1、2，报告 3 个种子的 均值 ± 标准差
    3. 长度外推：训练只见过 64 个字，测到 4096 个字
结果：experiments/compare.md（最新）+ experiments/results/<时间>-compare.md（存档）
"""
import sys

import numpy as np

from common import (CFG, ITERS, T, batches, evaluate, forward_flops_per_char, is_zeropoint, prepare_data,
                    save_result, slot_usage, train)
from models import TransformerLM, ZeroX
import zerocg as zc

LRS = [1e-3, 3e-3, 6e-3]
SEEDS = [0, 1, 2]
MODELS = {
    "Transformer": lambda V: TransformerLM(V, T, **CFG),
    "Transformer (sin)": lambda V: TransformerLM(V, T, pos="sin", **CFG),
    "ZeroPoint": lambda V: ZeroX(V, n_slot=16, **CFG),
}
# 长度外推的分段；Transformer 的注意力矩阵是 长度² 的，4096 字放不进内存，只测到 1024
SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)]
LONG_MAX = {"Transformer (sin)": 1024, "ZeroPoint": 4096}
LONG_N = 8   # 测试序列条数


def long_context_loss(model, val_data, max_len):
    x, y = next(batches(val_data, 1, max_len, 7, batch=LONG_N))
    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 = -np.take_along_axis(logp, y[..., None], -1)[..., 0]
    return np.array([nll[:, a:b].mean() if b <= max_len else np.nan for a, b in SEGMENTS])


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else None
    text, V, train_data, val_data = prepare_data(path)
    print(f"语料 {len(text):,} 字符，词表 {V}，每次训练 {ITERS} 步", flush=True)

    sweep = {name: {} for name in MODELS}      # name -> lr -> 验证 loss（种子 0）
    finals = {name: [] for name in MODELS}     # name -> [各种子的验证 loss]
    long_loss = {name: [] for name in LONG_MAX}
    usage = []

    def run(name, seed, lr):
        np.random.seed(seed)
        model = MODELS[name](V)
        val = train(name, model, train_data, val_data, seed, lr)
        print(f"  => {name} seed={seed} lr={lr:g}: val {val:.4f}", flush=True)
        return model, val

    for name in MODELS:
        best = None
        for lr in LRS:
            model, val = run(name, 0, lr)
            sweep[name][lr] = val
            if best is None or val < best[1]:
                best = (model, val, lr)
        best_lr = best[2]
        trained = [best[0]] + [run(name, s, best_lr)[0] for s in SEEDS[1:]]
        finals[name] = [best[1]] + [None] * (len(SEEDS) - 1)
        for k, model in enumerate(trained):
            if k > 0:
                finals[name][k] = evaluate(model, val_data)
            if name in LONG_MAX:
                long_loss[name].append(long_context_loss(model, val_data, LONG_MAX[name]))
            if is_zeropoint(model):
                usage.append(slot_usage(model))

    # ---- 写结果 ----
    m0 = {name: make(V) for name, make in MODELS.items()}
    lines = ["# ZeroPoint vs Transformer 对照实验\n",
             f"- 语料：{'Python 标准库源码' if not path else path}，{len(text):,} 字符，词表 {V}；90% 训练 / 10% 验证",
             f"- 三个模型：宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头、Pre-LN、GELU MLP 4 倍；"
             "ZeroPoint 每头 16 个槽、可学习初始状态、分块并行（块长 32）、无位置编码",
             f"- 训练：上下文 {T}、batch 32、{ITERS} 步、AdamW（权重衰减 0.01）、预热 100 步 + 余弦、梯度裁剪 1.0",
             f"- 学习率：每个模型在 {LRS} 里各试一次（种子 0），取最好的再跑种子 {SEEDS[1:]}",
             "- 验证 loss = 40 个固定验证 batch 的平均交叉熵（nats/字符，越低越好），均值 ± 标准差（3 个种子）",
             "- 计算量 = 前向时每个字的矩阵乘浮点运算次数（不受机器负载影响；训练一步约 3 倍）\n",
             "| 模型 | 参数量 | 最佳学习率 | 验证 loss | 前向计算量（每字） | 推理状态（第 64 字 / 第 4096 字，浮点数） |",
             "|---|---|---|---|---|---|"]
    for name, vals in finals.items():
        m = m0[name]
        best_lr = min(sweep[name], key=sweep[name].get)
        state = (f"{m.state_floats():,} / {m.state_floats():,}（固定）" if is_zeropoint(m)
                 else f"{m.state_floats(64):,} / {m.state_floats(4096):,}（KV 缓存）")
        lines.append(f"| {name} | {m.num_parameters():,} | {best_lr:g} | {np.mean(vals):.4f} ± {np.std(vals):.4f} | "
                     f"{forward_flops_per_char(m, V) / 1e3:.0f} K | {state} |")

    lines += ["", "## 学习率搜索（种子 0 的验证 loss）\n",
              "| 模型 | " + " | ".join(f"lr={lr:g}" for lr in LRS) + " |",
              "|---|" + "---|" * len(LRS)]
    for name, d in sweep.items():
        lines.append(f"| {name} | " + " | ".join(f"{d[lr]:.4f}" for lr in LRS) + " |")

    lines += ["", f"## 长度外推（训练上下文 {T}；{LONG_N} 条验证序列 × 3 个种子平均）\n",
              "| 位置 | ZeroPoint | Transformer (sin) | Transformer（可学习位置编码） |", "|---|---|---|---|"]
    lz, ls = np.nanmean(long_loss["ZeroPoint"], axis=0), np.nanmean(long_loss["Transformer (sin)"], axis=0)
    for k, (a, b) in enumerate(SEGMENTS):
        tag = "（训练见过的长度）" if b <= T else ""
        s = "—（注意力矩阵太大，没测）" if np.isnan(ls[k]) else f"{ls[k]:.4f}"
        lines.append(f"| {a} ~ {b - 1}{tag} | {lz[k]:.4f} | {s} | {'能跑' if b <= T else f'跑不了（只有 {T} 个位置）'} |")

    u = np.mean(usage, axis=0)   # (层, 2)
    lines += ["", "## ZeroPoint 槽的使用情况（有效槽数，满分 16）\n",
              "有效槽数 = 平均读/写地址分布的 exp(熵)：16 个槽用得一样多 = 16，全挤在一个槽 = 1。\n",
              "| 层 | 读 | 写 |", "|---|---|---|"]
    lines += [f"| {i} | {r:.1f} | {w:.1f} |" for i, (r, w) in enumerate(u)]
    save_result("compare", lines)


if __name__ == "__main__":
    main()