"""
=============================================================================
CS486 · 第 6 讲 · 代码 02：长度泛化——到底是哪一种"长"？
=============================================================================

【这个脚本要回答的问题】（讲义 §4）
    第 1、4、5 讲反复看到：cot_L2 在 n ≤ 8 的题上训练，n = 12、16 时大量失败（主要是"停不下来"）。
    "题目更长"其实混在一起的有好几件事：
        (a) 题面更长（陈述更多，上下文更长，位置更远）；
        (b) 环更大（要在更多的字母里找）；
        (c) 要走的步数更多（推理链更长）。
    要知道是哪一个，就要**分别改变**它们（控制变量，第 1 讲 §8.3）。我们训练一个对照模型 distract：
    和 cot_L2 完全相同（环大小 4–8、跳数 1 … n-1、同样的结构、步数、种子），只是每道训练题再加上
    若干条干扰陈述（另一个和答案无关的环），让题面最长有 16 条陈述——它见过长题面，但从没见过大于 8 的环。

    演示 1：两个模型在 n = 8 … 16 的环上的准确率。
    演示 2：环固定为 n = 8，只增加干扰陈述（题面变长、环不变）。
    演示 3：失败的样子。
    演示 4：比训练时更多的步数（k = 8、9）——一个混淆因素。

【运行】
    python3 02_length_generalization.py           # 第一次约 50 秒（训练 distract 模型），之后约 10 秒
    python3 02_length_generalization.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 reliability_lab.py。
"""

import argparse
import os
import random
import sys
import time
from collections import Counter

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

import hop_world as hw                                                   # noqa: E402
import reliability_lab as rlab                                           # noqa: E402
import rl_lab as rl                                                      # noqa: E402
import tiny_reasoner as tr                                               # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

NS = (8, 10, 12, 14, 16)
DS = (0, 2, 4, 6, 8)


def accuracy_and_kinds(model, probs):
    """贪心作答，返回 (准确率, 结果类型计数)。max_new = k + 5（和第 1、4、5 讲一致）。"""
    outs = tr.generate(model, [p.prompt for p in probs], max(p.k for p in probs) + 5)
    kinds = Counter(rl.classify(p, rl.response_of(o)) for p, o in zip(probs, outs))
    return kinds["correct"] / len(probs), kinds


def distract_set(d, per_k):
    """环大小 8、k = 3–6 各 per_k 道、外加 d 条干扰陈述（d = 0 时就是普通的 n = 8 测试集）。"""
    out = []
    for k in (3, 4, 5, 6):
        out += hw.make_eval_set(8, k, per_k, n_distract=d, seed=61)
    return out


def main():
    ap = argparse.ArgumentParser(description="长度泛化：是哪一种长？")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    per_k = 5 if args.quick else 50

    print_header("准备：cot_L2（第 1 讲）和 distract（同样的训练，只多了干扰陈述）")
    models = {"cot_L2": tr.get_model("cot_L2", quick=args.quick, verbose=not args.quick)[0],
              "distract": rlab.sft_model("distract", quick=args.quick)}
    rng = random.Random(0)
    ex = rlab._distract_pairs(rng, 3)
    print(f"  distract 的训练样本举例：{[a + ' → ' + b for a, b in ex[:2]]}")

    print_header("演示 1：改变环的大小 n（k = 3–6 各 %d 道；两个模型都只见过 n ≤ 8 的环）" % per_k)
    acc1 = {}
    rows = []
    for name, model in models.items():
        cells = [name]
        for n in NS:
            a, _ = accuracy_and_kinds(model, rl.eval_set(n, per_k=per_k))
            acc1[(name, n)] = a
            cells.append(pct(a))
        rows.append(cells)
    print_table(["模型", *[f"n = {n}" for n in NS]], rows, [10] + [9] * len(NS))
    print(f"  （题面长度：n = 8 时约 {len(rl.eval_set(8, per_k=1)[0].prompt)} 个字符，n = 16 时约 "
          f"{len(rl.eval_set(16, per_k=1)[0].prompt)} 个字符）")

    print_header("演示 2：环固定为 n = 8，只加干扰陈述（题面变长，要走的路一点没变）")
    acc2, rows = {}, []
    for name, model in models.items():
        cells = [name]
        for d in DS:
            a, _ = accuracy_and_kinds(model, distract_set(d, per_k))
            acc2[(name, d)] = a
            cells.append(pct(a))
        rows.append(cells)
    print_table(["模型", *[f"干扰 {d} 条" for d in DS]], rows, [10] + [10] * len(DS))
    print(f"  （干扰 8 条时题面有 16 条陈述，和演示 1 的 n = 16 一样长）")

    print_header("演示 3：失败的样子——cot_L2 在两种'长题面'上的结果类型")
    rows, kinds3 = [], {}
    for label, probs in (("n = 16（大环）", rl.eval_set(16, per_k=per_k)),
                         ("n = 8 + 干扰 8 条", distract_set(8, per_k))):
        _, c = accuracy_and_kinds(models["cot_L2"], probs)
        kinds3[label] = c
        outs = tr.generate(models["cot_L2"], [probs[0].prompt], probs[0].k + 5)
        rows.append([label] + [c.get(t, 0) for t in rl.FAILURE_TYPES] + [f"{probs[0].trace} → {rl.response_of(outs[0])}"])
    print_table(["测试集", "对", "走过头", "停早了", "查找错", "抄错", "没写 |", "例：正确推理链 → 模型写的"],
                rows, [18, 5, 7, 7, 7, 5, 7, 30])

    print_header("演示 4：比训练时更多的步数（n = 12，每个 k 各 %d 道）" % (per_k * 2))
    rows, kinds4 = [], {}
    for name, model in models.items():
        for k in (6, 7, 8, 9):
            probs = hw.make_eval_set(12, k, per_k * 2, seed=21)
            a, c = accuracy_and_kinds(model, probs)
            kinds4[(name, k)] = (a, c)
            outs = tr.generate(model, [probs[0].prompt], k + 5)
            rows.append([name, k, pct(a), c.get("early", 0), c.get("overshoot", 0) + c.get("malformed", 0),
                         f"{probs[0].trace} → {rl.response_of(outs[0])}"])
    print_table(["模型", "k", "准确率", "停早了", "停不下来", "例：正确推理链 → 模型写的"], rows, [10, 4, 8, 8, 10, 32])
    seen = sorted({c for p in [hw.sample_train_problem(random.Random(s), (4, 8)) for s in range(2000)] for c in str(p.k)})
    print(f"  训练分布里题面上出现过的跳数数字：{''.join(seen)}——数字 8、9 从没出现过。")

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    c16, c8d = kinds3["n = 16（大环）"], kinds3["n = 8 + 干扰 8 条"]
    print_observations([
        f"只加干扰陈述（演示 2），cot_L2 的准确率就从 {pct(acc2[('cot_L2', 0)])} 一路掉到 {pct(acc2[('cot_L2', 8)])}"
        f"（干扰 2 / 4 / 6 / 8 条：{' / '.join(pct(acc2[('cot_L2', d)]) for d in DS[1:])}），"
        f"和演示 1 里把环从 8 变成 10 / 12 / 14 / 16 的曲线（{' / '.join(pct(acc1[('cot_L2', n)]) for n in NS[1:])}）几乎重合。"
        "要走的路一点都没变，变的只是题面的长度——**cot_L2 的'长度泛化失败'是题面（上下文）长度的问题，不是环的大小、也不是推理的步数**。",
        f"失败的样子也一样（演示 3）：n = 16 和'n = 8 + 8 条干扰'上，没写 '|' 的分别是 {c16['malformed']} 和 {c8d['malformed']} 道、"
        f"走过头 {c16['overshoot']} 和 {c8d['overshoot']} 道、查找错 {c16['lookup']} 和 {c8d['lookup']} 道。"
        "两种'长'在模型看来是同一件事：题面一长，它就判断不准该在哪里停下。",
        f"所以修补很便宜（演示 1、2）：训练时加上干扰陈述，让题面长度覆盖到 16 条陈述，distract 在 n = 10 … 16 的大环上"
        f"全部 {pct(acc1[('distract', 16)])}——它从没见过大于 8 的环。它学到的是'在长上下文里照样查找、照样数步数'，"
        "大环对它来说只是另一种长题面。这和真实 LLM 的长上下文训练是同一个道理：**泛化到更长的输入，往往需要训练分布里有足够长的输入**，"
        "而不需要更长的'推理'。",
        f"推理步数是另一根轴（演示 4）：distract 在 k = 6、7 时 {pct(kinds4[('distract', 7)][0])}，k = 8、9 时全部停早了"
        f"（{kinds4[('distract', 8)][1].get('early', 0)} / {kinds4[('distract', 9)][1].get('early', 0)} 道）。"
        "但这里有一个混淆：训练题的跳数最多是 7，题面里的数字 8、9 **从来没出现过**，模型不知道它们是什么意思"
        "（第 1 讲 §4.4 也遇到过）。这个失败说明不了'模型不能做更多步的推理'——要检验步数泛化，"
        "得换一种不引入新符号的设计（例如第 3 讲的循环模型在测试时多循环几轮）。",
        "方法上的教训：一个'泛化失败'通常同时改变了好几个因素；在下结论之前，构造只改变其中一个因素的测试集。"
        "代码 03 会看到，在复杂度扫描实验里不这样做会得出什么样的结论。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()