"""
=============================================================================
CS486 · 第 3 讲 · 代码 03：不写 token 的推理——循环 Transformer 与"多想几轮"
=============================================================================

【这个脚本要回答的问题】（讲义 §6）
    第 1 讲说思维链有用，是因为"写下的 token 从第 0 层重新进入模型"带来了串行深度。
    那么串行深度一定要靠写 token 吗？我们训练一个**循环 Transformer**：
        一个 2 层的循环体（参数量和第 1 讲的 2 层模型完全相同：104,320），在同一道题上重复执行 T 次，
        中间结果留在隐藏状态里，**不写任何 token**；第 T 轮之后在起点位置读出答案。
    题面去掉了跳数数字（"…;?a"）：走几跳只由循环次数 T 决定。训练时 k ∈ [1, 4]，T = k。

    三个版本：
        逐轮监督 + 输入注入：每一轮的读出都要预测"走了 t 跳之后的字母"（中间监督，像思维链），
                              每轮开始时把原始词嵌入加回去
        逐轮监督，不注入
        只监督最终答案 + 输入注入：像直接回答，只是模型"更深"

    演示 1：k = 1–4（训练范围）和 k = 5–8（更长的题，测试时多循环几轮）的准确率。
    演示 2：循环次数和跳数不匹配时——少想几轮、多想几轮，模型会停在哪里？
    演示 3：和第 1 讲的模型并排比较（同样的参数量）。

【运行】
    python3 03_looped_latent.py           # 第一次约 1.5 分钟（并行训练 3 个循环模型），之后约 5 秒
    python3 03_looped_latent.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 cot_lab.py + 第 1、2 讲的模块。
"""

import argparse
import os
import sys
import time

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

import cot_lab as cl                                                     # noqa: E402
import hop_world as hw                                                   # noqa: E402
import tiny_reasoner as tr                                               # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

BASE = {"kind": "looped", "steps": 2000, "kmax": 4, "seed": 0}
JOBS = [
    dict(BASE, supervision="every", inject=True),
    dict(BASE, supervision="every", inject=False),
    dict(BASE, supervision="final", inject=True),
]
LABELS = ["逐轮监督 + 输入注入", "逐轮监督，不注入", "只监督最终答案 + 注入"]


def accuracy_at(model, k, T, count, seed=11):
    """
    循环 T 轮，在 n = 8、k 跳的题上，读出第 T 轮的答案，算它等于"走 k 跳"的比例。

    【返回】(准确率, 第 T 轮读出等于"走 T 跳的位置"的比例)
    """
    probs = hw.make_eval_set(8, k, count, seed=seed)
    preds = cl.looped_predict(model, probs, T)
    acc = sum(pr[-1] == p.answer for pr, p in zip(preds, probs)) / count
    succ = [{st[0]: st[2] for st in p.statements} for p in probs]
    at_T = 0
    for pr, p, s in zip(preds, probs, succ):
        x = p.start
        for _ in range(T):
            x = s[x]
        at_T += pr[-1] == x
    return acc, at_T / count


def main():
    ap = argparse.ArgumentParser(description="循环 Transformer 与潜在推理")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    count = 30 if args.quick else 300

    print_header("准备：三个循环模型（2 层循环体，权重共享；训练时 k ∈ [1, 4]、循环 k 轮；2000 步）")
    p = hw.make_eval_set(8, 3, 1, seed=11)[0]
    print(f"  题面（没有跳数数字）：{cl.latent_prompt(p)}    循环 3 轮后应读出 {p.answer}")
    models = cl.load_jobs(JOBS, quick=args.quick)
    print(f"  参数量：循环模型 {sum(q.numel() for q in models[0].parameters()):,}；"
          f"第 1 讲的 2 层模型 {tr.count_params(tr.TinyReasoner(tr.ModelConfig())):,}")

    print_header(f"演示 1：循环 k 轮，回答 k 跳的题（n = 8，每格 {count} 道题；k ≥ 5 是训练时没见过的长度）")
    res = {lab: [accuracy_at(m, k, k, count)[0] for k in range(1, 9)] for lab, m in zip(LABELS, models)}
    rows = [[f"k = {k}" + ("（外推）" if k > 4 else "")] + [pct(res[lab][k - 1]) for lab in LABELS] for k in range(1, 9)]
    print_table(["跳数（= 循环轮数）"] + LABELS, rows, [18, 22, 20, 24])
    print("（k = 8 时 n = 8 的环走回起点；训练时 k ≤ 4 < n，答案从来不是起点。）")

    print_header("演示 2：循环轮数 T 和题目跳数 k 不一致时（逐轮监督 + 输入注入的模型，k = 4）")
    m0 = models[0]
    rows, mismatch = [], {}
    for T in (1, 2, 3, 4, 5, 6, 8):
        acc, at_T = accuracy_at(m0, 4, T, count)
        mismatch[T] = (acc, at_T)
        rows.append([T, pct(acc), pct(at_T)])
    print_table(["循环轮数 T", "读出 = 走 4 跳（正确答案）", "读出 = 走 T 跳的位置"], rows, [12, 26, 22])

    print_header("演示 3：同样的参数量，四种'想'的方式（n = 8）")
    direct = tr.get_model("direct_L2", quick=args.quick)[0]
    cot = tr.get_model("cot_L2", quick=args.quick)[0]
    rows = []
    for k in (1, 2, 4, 7):
        probs = hw.make_eval_set(8, k, count, seed=11)
        a_dir = tr.accuracy(tr.evaluate(direct, probs, "direct"))
        a_cot = tr.accuracy(tr.evaluate(cot, probs, "cot"))
        rows.append([f"k = {k}", pct(a_dir), pct(a_cot), pct(res[LABELS[0]][k - 1]), pct(res[LABELS[2]][k - 1])])
    print_table(["跳数", "直接回答（第 1 讲）", "思维链（第 1 讲）", "循环，逐轮监督", "循环，只监督答案"], rows,
                [8, 20, 18, 16, 18])
    print("（直接回答和思维链训练时 k ≤ 7；循环模型训练时 k ≤ 4，k = 7 对它是外推。）")

    if args.quick:
        print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    a, b, c = (res[lab] for lab in LABELS)
    in_range = "全部 100.0%" if min(a[:4]) == max(a[:4]) == 1.0 else f"{pct(min(a[:4]))}–{pct(max(a[:4]))}"
    print_observations([
        f"逐轮监督的循环模型在训练范围 k = 1–4 上{in_range}——同样 104,320 个参数、"
        "不写任何 token，只是把 2 层循环 k 次，它就做到了第 1 讲的 2 层直接回答模型做不到的多跳推理。"
        "串行深度可以来自'写 token'，也可以来自'循环'。",
        f"测试时多循环几轮，它能解出比训练时更长的题：k = 5、6、7 时是 {pct(a[4])}、{pct(a[5])}、{pct(a[6])}，"
        "外推随长度逐渐衰减。这就是'测试时多想一会儿'在潜在空间里的版本。"
        f"k = 8 时只有 {pct(a[7])}：n = 8 的环走 8 跳回到起点，而训练时答案从来不是起点——这一格混进了别的因素。",
        f"输入注入的作用在这里**不稳健**：不注入的版本在 k = 5、6、7、8 上是 {pct(b[4])}、{pct(b[5])}、{pct(b[6])}、{pct(b[7])}，"
        "在长的几格上反而更好。编写时的原型（题面里还带着跳数数字）得到的是相反的结论（注入明显更好）。"
        "一个种子、一次训练的比较，不足以下'注入有用 / 没用'的结论（第 1 讲 §8.3、第 2 讲 §3.3）。",
        f"只监督最终答案的循环模型在所有 k 上都在随机水平（k = 1–4：{pct(min(c[:4]))}–{pct(max(c[:4]))}）。"
        "和第 1 讲一样：结构上的深度不够，还需要**中间监督**才学得会。循环提供了表达能力，逐轮监督提供了可学习性。",
        f"演示 2：k = 4 的题，循环 T 轮时读出几乎总是'走 T 跳的位置'（T = 2：{pct(mismatch[2][1])}；T = 5："
        f"{pct(mismatch[5][1])}；T = 6：{pct(mismatch[6][1])}），只有 T = 4 时才是正确答案（{pct(mismatch[4][0])}）。"
        "模型每循环一轮就走一跳，**它自己不知道什么时候该停**——停止条件必须由外部给出（这里是 T）。"
        "真实的推理模型要自己决定想多久，这是自适应计算（adaptive computation）要解决的问题（讲义 §6.4）。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()