"""
=============================================================================
CS486 · 第 6 讲 · 代码 03：一个"推理崩溃"是怎样被实验设计制造出来的——复盘"思考的幻觉"之争
=============================================================================

【这个脚本要回答的问题】（讲义 §6）
    2025 年 6 月，Shojaee et al.（Apple，"The Illusion of Thinking"）用可调复杂度的谜题扫描推理模型，
    报告：复杂度超过某个阈值后，准确率**崩溃到 0**；而且在崩溃附近，模型用的推理 token 反而**变少**。
    很快就有反驳（Lawsen 2025，"The Illusion of the Illusion of Thinking"）指出几个实验设计问题：
        (1) 输出长度上限：汉诺塔的完整解法有 2^N − 1 步，N 大时写完整个解法超出了模型的输出预算；
        (2) 无解的实例：一部分过河问题在给定的船只容量下根本无解，却被计为模型失败；
        (3) 评分方式：要求写出完整的动作序列，而不是允许更紧凑的表示。
    这场争论的教训不在于谁对谁错，而在于：**一条"随复杂度崩溃"的曲线，可能来自模型，也可能来自实验设计**。

    我们在实验台上把这件事做一遍。"推理模型"是代码 02 的 distract（在 n ≤ 16、k ≤ 7 的题上几乎完美），
    复杂度 = 跳数 k（1 … 9），环大小 n = 16。一个"朴素"的基准有三个缺陷，都是可以控制的开关：
        预算  ：输出最多 8 个 token（k ≥ 7 时写不完"推理链 + | + 答案"）；
        无解题：k ≥ 5 时，生成器以 40% 的概率删掉路径上的一条陈述（题目无解），但仍按原答案判分；
        新符号：k = 8、9 的题面数字在训练分布里从没出现过（代码 02 演示 4）。

    演示 1：朴素的扫描——准确率和"推理长度"随复杂度的变化。
    演示 2：逐个去掉缺陷，看"崩溃"还剩多少。
    演示 3：最后剩下的那部分——换一个训练分布覆盖 k = 8、9 的模型（distract9）。
    演示 4：模型在无解题上写了什么？

【运行】
    python3 03_illusion_of_thinking.py           # 第一次约 1 分钟（训练 distract9；distract 读代码 02 的缓存），之后约 10 秒
    python3 03_illusion_of_thinking.py --quick   # 冒烟测试

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

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 reliability_lab as rlab                                           # noqa: E402
import tiny_reasoner as tr                                               # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

KS = tuple(range(1, 10))
NAIVE_BUDGET = 8
DROP_FROM, DROP_PROB = 5, 0.4


def sweep(model, count, budget, drop, adequate=False):
    """
    在 k = 1 … 9 上各做 count 道题。

    【参数】budget：固定的输出预算（adequate=True 时改用 k + 5，足够写完）；drop：是否打开"无解题"缺陷
    【返回】{k: run_sweep 的结果列表}
    """
    out = {}
    for k in KS:
        items = rlab.sweep_problems(k, count, drop_from=DROP_FROM if drop else None, drop_prob=DROP_PROB)
        out[k] = rlab.run_sweep(model, items, k + 5 if adequate else budget)
    return out


def acc(res, solvable_only=False):
    """准确率；solvable_only=True 时只在可解的题上算。【边界】没有题时返回 None。"""
    rs = [r for r in res if r["solvable"] or not solvable_only]
    return sum(r["correct"] for r in rs) / len(rs) if rs else None


def main():
    ap = argparse.ArgumentParser(description="复杂度扫描的设计陷阱")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    count = 10 if args.quick else 200

    print_header("准备：distract（代码 02）和 distract9（训练时跳数覆盖 1–9）")
    model = rlab.sft_model("distract", quick=args.quick)
    model9 = rlab.sft_model("distract9", quick=args.quick)

    print_header(f"演示 1：朴素的扫描（每个 k {count} 道题；输出预算 {NAIVE_BUDGET} 个 token；"
                 f"k ≥ {DROP_FROM} 时 {pct(DROP_PROB, 0)} 的题无解）")
    naive = sweep(model, count, NAIVE_BUDGET, drop=True)
    rows = []
    for k in KS:
        res = naive[k]
        rows.append([k, pct(acc(res)), f"{sum(r['trace_len'] for r in res) / len(res):.2f}",
                     pct(sum(r["truncated"] for r in res) / len(res), 0)])
    print_table(["复杂度 k", "准确率", "平均推理长度", "被截断"], rows, [10, 10, 14, 10])
    print("  （平均推理长度：'|' 之前写了几个字母；被截断：预算用完时还没写出答案字母）")

    print_header("演示 2：逐个去掉缺陷")
    fixed_budget = sweep(model, count, NAIVE_BUDGET, drop=True, adequate=True)
    clean = sweep(model, count, NAIVE_BUDGET, drop=False, adequate=True)
    lines = [("朴素（预算 8 + 无解题）", naive, False), ("+ 足够的预算（k + 5）", fixed_budget, False),
             ("+ 只在可解的题上算", fixed_budget, True), ("+ 生成器修好（全部可解）", clean, False)]
    rows, table = [], {}
    for label, res, solv in lines:
        vals = [acc(res[k], solv) for k in KS]
        table[label] = vals
        rows.append([label] + [pct(v, 0) for v in vals])
    print_table(["设计", *[f"k={k}" for k in KS]], rows, [24] + [6] * len(KS))

    print_header("演示 3：剩下的 k = 8、9——换一个训练时见过这两个数字的模型")
    res9 = sweep(model9, count, NAIVE_BUDGET, drop=False, adequate=True)
    rows = [["distract（训练 k ≤ 7）"] + [pct(acc(clean[k]), 0) for k in KS],
            ["distract9（训练 k ≤ 9）"] + [pct(acc(res9[k]), 0) for k in KS]]
    print_table(["模型", *[f"k={k}" for k in KS]], rows, [24] + [6] * len(KS))
    lens = {k: sum(r["trace_len"] for r in clean[k]) / len(clean[k]) for k in (7, 8, 9)}
    print(f"  distract 在 k = 7 / 8 / 9 时平均写 {lens[7]:.2f} / {lens[8]:.2f} / {lens[9]:.2f} 步就停下")

    print_header("演示 4：模型在无解题上写了什么？（k = 6，distract，足够的预算）")
    items = rlab.sweep_problems(6, count, drop_from=DROP_FROM, drop_prob=DROP_PROB)
    bad = [(p, ok) for p, ok in items if not ok][:4]
    raw = tr.generate(model, [p.prompt for p, _ in bad], 11)
    for (p, _), o in zip(bad, raw):
        missing = next(c for c in p.walk[:-1] if c not in {st[0] for st in p.statements})
        print(f"  题面缺了 {missing!r} 的出边；生成器的'标准答案' {p.answer!r}；模型写的 {o.split('$')[0] + '$'!r}")
    n_bad = sum(1 for _, ok in items if not ok)
    print(f"  （k = 6 的 {count} 道题里 {n_bad} 道无解；模型没有'无解'这个选项，总会写出某个答案）")

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    nv = {k: acc(naive[k]) for k in KS}
    ln = {k: sum(r["trace_len"] for r in naive[k]) / len(naive[k]) for k in KS}
    print_observations([
        f"朴素的扫描（演示 1）画出了一条典型的'推理崩溃'曲线：k ≤ 4 时 {pct(nv[4])}，k = 5、6 掉到 {pct(nv[5])} / {pct(nv[6])}，"
        f"k = 7 是 {pct(nv[7])}，k = 8、9 只有 {pct(nv[8])} / {pct(nv[9])}。而且'推理长度'在崩溃处**变短**了："
        f"k = 7 时平均写 {ln[7]:.2f} 步，k = 8、9 时只写 {ln[8]:.2f} / {ln[9]:.2f} 步——看上去就像'问题太难时，模型反而少想了'。"
        "如果只看这张表，很容易写出'推理模型在复杂度 5 以上崩溃，并且在崩溃时减少推理努力'这样的结论。",
        f"逐个去掉缺陷（演示 2）：给足输出预算，k = 7 从 {pct(table['朴素（预算 8 + 无解题）'][6], 0)} 回到 "
        f"{pct(table['+ 足够的预算（k + 5）'][6], 0)}（k = 7 的推理链 + '|' + 答案需要 9 个 token，预算只有 8 个——"
        "模型写到了最后一步，却没地方写答案）；只在可解的题上算，k = 5–7 全部回到 100%（之前的 36–40% 失败全是无解题）。"
        "这两个缺陷正对应 Lawsen（2025）对汉诺塔（输出长度）和过河问题（无解实例）的批评。",
        f"剩下的 k = 8、9（演示 3）在修好生成器之后仍然是 {pct(acc(clean[8]), 0)}：distract 总是走 5 步就停下。"
        f"换一个训练时见过数字 8、9 的模型（distract9，其他完全相同），k = 8、9 是 {pct(acc(res9[8]), 0)} / {pct(acc(res9[9]), 0)}。"
        "所以最后这段'崩溃'也不是'推理能力的上限'，而是**训练分布的边界**（题面里出现了没见过的符号）。"
        "这一条比前两条更隐蔽：它确实是模型的失败，但对它的正确描述是'不能泛化到没见过的步数符号'，而不是'不能做长推理'。",
        f"无解的题（演示 4）上，模型不会说'无解'——它没有这个选项——而是在缺失的那一步**编造**一个转移，写出一条看上去完整的推理链"
        "（例如 'm' 没有出边，它照样写 'mkhbol|l$'）。真实的推理模型也一样：面对无解或题意有误的问题，它们常常给出一个自信的答案。"
        "一个好的基准要么保证每道题可解，要么允许并奖励'无解'的回答。",
        "实验设计的检查清单（讲义 §6.4）：输出预算够不够写完一个完整的解？每个实例都可解吗（用独立的求解器验证）？"
        "失败里有多少是截断、格式、无解，有多少是真正答错？复杂度的增加有没有同时引入别的变化（新符号、更长的输入）？"
        "把这些都控制住之后，剩下的曲线才是关于模型的。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()