"""
=============================================================================
CS486 · 第 1 讲 · 代码 03：是在推理，还是在走捷径？——用"扰动探针"检验一个模型学到了什么
=============================================================================

【这个脚本要回答的问题】（讲义 §4）
    一个模型在测试集上 100% 正确，能说明它"学会了推理"吗？不能。
    测试集和训练集来自同一个分布时，一个**捷径**（shortcut）——只在这个分布上碰巧有效的规则——
    也能拿满分。要区分"学会了算法"和"学会了捷径"，研究者会构造**扰动探针**（perturbation probes）：
    保持题目的本质不变，只改变"不该影响答案"的表面特征，看准确率会不会掉
    （GSM-Symbolic 换数字和名字、GSM-NoOp 加无关句子、反事实任务……讲义 §4.2）。

    我们训练两个思维链模型，结构和步数完全相同，只有训练题的**陈述顺序**不同：
        cot_L2        ：陈述随机打乱（代码 01 的模型）
        cot_sorted_L2 ：陈述按行走顺序排列（a>c; c>f; f>b; …）——此时"第 i 条陈述的箭头右边
                        就是第 i 跳的答案"，一个**按位置读**的捷径就能拿满分。

    演示 1：两个模型在各自的训练分布上都接近满分——只看分布内的准确率，分不出谁学到了什么。
    演示 2：顺序探针——打乱 / 按顺序 / 倒序。
    演示 3：诊断捷径——sorted 模型在打乱的题上，写出的推理字母符合"按名字查找"还是"按位置读"？
    演示 4：干扰陈述 vs 题面长度——一个**控制变量**的设计：加干扰陈述准确率掉了，
            是因为"被无关信息干扰"，还是仅仅因为"题面变长了"？
    演示 5：比训练时更多的跳数（题面长度不变）——推理链能不能比训练时更长？

【运行】
    python3 03_reasoning_or_shortcut.py           # 第一次约 1 分钟（训练 cot_sorted_L2；若没有缓存还要训练 cot_L2）
    python3 03_reasoning_or_shortcut.py --quick   # 冒烟测试，几秒钟，数字没有意义

【依赖】PyTorch + 本目录的 hop_world.py、tiny_reasoner.py、research_kit.py、console.py。
"""

import argparse
import time
from collections import Counter

import hop_world as hw
import tiny_reasoner as tr
from console import pct, print_header, print_observations, print_table
from research_kit import wilson_interval

MODELS = (("cot_L2", "打乱顺序训练"), ("cot_sorted_L2", "按顺序训练"))


def acc_on(model, n, k, count, order="shuffle", n_distract=0, seed=5):
    """在一个 (n, k, order, n_distract) 评测集上的准确率与逐题结果。【返回】(acc, probs, results)"""
    probs = hw.make_eval_set(n, k, count, order=order, n_distract=n_distract, seed=seed)
    res = tr.evaluate(model, probs, "cot")
    return tr.accuracy(res), probs, res


def with_ci(acc, count):
    """把准确率写成 "xx.x% [lo, hi]"（Wilson 95% 区间）。"""
    lo, hi = wilson_interval(round(acc * count), count)
    return f"{pct(acc)} [{pct(lo, 0)}, {pct(hi, 0)}]"


def demo1_in_distribution(models, count):
    print_header("演示 1：分布内的准确率（各自的训练分布：n = 8，k = 1–7）")
    rows, out = [], {}
    for (name, label), order in zip(MODELS, ("shuffle", "sorted")):
        accs = [acc_on(models[name], 8, k, count // 2, order=order)[0] for k in range(1, 8)]
        out[name] = sum(accs) / len(accs)
        rows.append([f"{name}（{label}）", order, pct(out[name]), pct(min(accs))])
    print_table(["模型", "测试顺序", "平均准确率", "最低的 k"], rows, [28, 10, 12, 10], ["<", "<", ">", ">"])
    return out


def demo2_order(models, count):
    print_header("演示 2：顺序探针（n = 8，k ∈ {2, 4, 7}，每种顺序共 3 × " + str(count) + " 道题）")
    rows, out = [], {}
    for name, label in MODELS:
        row = [f"{name}（{label}）"]
        for order in hw.ORDERS:
            c = n = 0
            for k in (2, 4, 7):
                a, probs, _ = acc_on(models[name], 8, k, count, order=order)
                c += round(a * len(probs))
                n += len(probs)
            out[(name, order)] = c / n
            row.append(with_ci(c / n, n))
        rows.append(row)
    print_table(["模型", "打乱 shuffle", "按顺序 sorted", "倒序 reversed"], rows, [28, 22, 22, 22])
    return out


RULES = (("name", "按名字查找（正确）"),
         ("absolute", "第 i 条陈述的右边"),
         ("next_src", "邻居规则 A"),
         ("next_tgt", "邻居规则 B"),
         ("either", "A 或 B"))


def rule_predictions(prob, step, cur):
    """
    推理链第 step 步（从 0 数起）、当前字母为 cur 时，几条候选规则各自会写出什么字母。

    【思路】在"按顺序排列"的训练题里，下面几条规则给出的答案**完全相同**，所以训练数据无法区分它们：
        name     ：按名字查找——找到以 cur 开头的陈述 "cur>x"，写 x（正确的、能泛化的规则）
        absolute ：按绝对位置——写第 step 条陈述的右边（不看 cur）
        next_src ：邻居规则 A——找到以 cur 开头的陈述，写**下一条陈述的左边**
                   （按顺序排列时，"a>c;c>f"里 c>f 之后的陈述以 f 开头，正好是答案）
        next_tgt ：邻居规则 B——找到以 cur 结尾的陈述，写**下一条陈述的右边**
                   （按顺序排列时，"a>c;c>f" 里 a>c 的下一条是 c>f，右边 f 正好是答案）
        either   ：A 或 B 之一
    在打乱顺序的题上，它们给出不同的字母，于是可以看出模型实际在用哪一条。
    这种"训练数据允许很多种解释、模型挑了其中一种"的现象叫**欠定**（underspecification，D'Amour et al. 2020）。

    【参数】prob：HopProblem；step：第几步；cur：当前字母（用模型自己上一步写的字母，这样检验的是"每一步的规则"）
    【返回】dict：规则名 → 预测的字母（规则不适用时为 None，例如 cur 所在的陈述已经是最后一条）
    """
    sts = prob.statements
    src = next((j for j, st in enumerate(sts) if st[0] == cur), None)
    tgt = next((j for j, st in enumerate(sts) if st[2] == cur), None)
    pred = {
        "name": sts[src][2] if src is not None else None,
        "absolute": sts[step][2] if step < len(sts) else None,
        "next_src": sts[src + 1][0] if src is not None and src + 1 < len(sts) else None,
        "next_tgt": sts[tgt + 1][2] if tgt is not None and tgt + 1 < len(sts) else None,
    }
    return pred


def demo3_diagnose(models, count):
    print_header("演示 3：诊断——在打乱顺序的题上，推理链的每一步符合哪条规则？（n = 8，k = 7）")
    rows, out = [], {}
    for name, label in MODELS:
        probs = hw.make_eval_set(8, 7, count, order="shuffle", seed=6)
        res = tr.evaluate(models[name], probs, "cot")
        hits = {r: 0 for r, _ in RULES}
        total = 0
        for p, r in zip(probs, res):
            g = r["generated"]
            for i in range(p.k):
                cur = p.start if i == 0 else g[i - 1:i]
                if cur not in hw.LETTERS:
                    break
                pred = rule_predictions(p, i, cur)
                ch = g[i:i + 1]
                total += 1
                for rule in ("name", "absolute", "next_src", "next_tgt"):
                    hits[rule] += ch == pred[rule]
                hits["either"] += ch in (pred["next_src"], pred["next_tgt"])
        out[name] = {rule: hits[rule] / total for rule, _ in RULES}
        rows.append([f"{name}（{label}）"] + [pct(out[name][rule]) for rule, _ in RULES])
    print_table(["模型"] + [lab for _, lab in RULES], rows, [28, 18, 16, 10, 10, 8])
    print("（每一步用模型自己上一步写的字母作为'当前字母'；不同规则偶尔给出同一个字母，所以各列可以相加超过 100%。"
          "\n  在按顺序排列的题上，这 4 条规则给出的答案完全相同。）")
    p = hw.make_eval_set(8, 7, 1, order="shuffle", seed=6)[0]
    print(f"\n例：题面 {p.prompt}    正确推理链 {p.trace}")
    for name, label in MODELS:
        g = tr.evaluate(models[name], [p], "cot")[0]["generated"]
        print(f"    {name} 写出 {g.split('$')[0]}")
    return out


def demo4_distractors(model, count):
    print_header("演示 4：干扰陈述 vs 题面长度（cot_L2，k = 3；训练时的题面最多 8 条陈述、没有干扰）")
    conds = [(4, 0), (8, 0), (4, 4), (12, 0), (8, 4), (4, 8)]
    rows, out = [], {}
    for n, nd in conds:
        a, probs, _ = acc_on(model, n, 3, count, n_distract=nd, seed=7)
        out[(n, nd)] = a
        rows.append([n, nd, n + nd, len(probs[0].prompt), with_ci(a, len(probs))])
    print_table(["环大小 n", "干扰陈述", "陈述总数", "题面长度", "准确率 [95% 区间]"], rows, [10, 10, 10, 10, 22])
    return out


def demo5_more_hops(model, count):
    print_header("演示 5：超出训练范围的跳数（cot_L2；训练时 k ≤ n-1 ≤ 7，所以数字 8、9 从没在训练题里出现过）")
    rows, out = [], {}
    for n, k in ((5, 6), (4, 7), (8, 8), (8, 9), (5, 9)):
        a, probs, res = acc_on(model, n, k, count, seed=8)
        chains = [r["generated"].split(hw.ANSWER_MARK)[0] for r in res]
        chain_ok = sum(c == p.trace for c, p in zip(chains, probs)) / len(probs)
        prefix_ok = sum(c == p.trace[:len(c)] for c, p in zip(chains, probs)) / len(probs)
        lengths = Counter(len(c) for c in chains)
        common_len, share = lengths.most_common(1)[0]
        out[(n, k)] = {"acc": a, "chain_ok": chain_ok, "prefix_ok": prefix_ok, "len": common_len,
                       "len_share": share / len(probs)}
        note = "绕环，数字见过" if k <= 7 else "数字没见过"
        rows.append([f"n = {n}, k = {k}", note, pct(a), pct(chain_ok), pct(prefix_ok),
                     f"{common_len}（{pct(share / len(probs), 0)}）"])
    print_table(["题目", "说明", "答案准确率", "推理链全对", "写出的部分都对", "推理链长度（最常见）"], rows,
                [14, 14, 10, 10, 14, 20], ["<", "<", ">", ">", ">", ">"])
    return out


def main():
    ap = argparse.ArgumentParser(description="是在推理，还是在走捷径？")
    ap.add_argument("--quick", action="store_true", help="冒烟测试：每个模型只训练 60 步")
    args = ap.parse_args()
    t0 = time.perf_counter()
    count = 30 if args.quick else 300

    print_header("准备：两个思维链模型（2 层，2000 步），只有训练题的陈述顺序不同")
    models = {name: tr.get_model(name, quick=args.quick)[0] for name, _ in MODELS}

    ind = demo1_in_distribution(models, count)
    order = demo2_order(models, count // 3 if args.quick else 200)
    diag = demo3_diagnose(models, count)
    dis = demo4_distractors(models["cot_L2"], count)
    hops = demo5_more_hops(models["cot_L2"], count)

    if args.quick:
        print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    print_observations([
        f"分布内两个模型都接近满分（打乱训练 {pct(ind['cot_L2'])}，按顺序训练 {pct(ind['cot_sorted_L2'])}）。"
        "只看这一列，你会以为两个模型学到了同样的东西。",
        f"顺序探针把它们分开了：按顺序训练的模型在打乱的题上只有 {pct(order[('cot_sorted_L2', 'shuffle')])}，"
        f"在倒序的题上 {pct(order[('cot_sorted_L2', 'reversed')])}；打乱训练的模型在三种顺序上都是 "
        f"{pct(min(order[('cot_L2', o)] for o in hw.ORDERS))} 以上。",
        f"演示 3 找到了捷径——但不是我们最先猜的那一个：按顺序训练的模型在打乱的题上，每一步只有 "
        f"{pct(diag['cot_sorted_L2']['name'])} 符合按名字查找、{pct(diag['cot_sorted_L2']['absolute'])} 符合'第 i 条陈述的右边'"
        f"（都接近随机），却有 {pct(diag['cot_sorted_L2']['either'])} 符合两条'邻居规则'之一：它学会的是"
        "'找到当前字母所在的陈述，再看相邻的那一条'。在按顺序排列的训练题里，这几条规则和正确规则给出完全相同的答案，"
        "训练数据无法区分它们（欠定）；模型挑了对 RoPE 注意力来说最容易的'相对位置'规则。",
        f"演示 4：加 4 条干扰陈述、但总长度仍是 8 条（n = 4 + 4）时准确率 {pct(dis[(4, 4)])}，"
        f"和同样 8 条、没有干扰的 n = 8（{pct(dis[(8, 0)])}）相近；总长度变成 12 条时，"
        f"无论有没有干扰都明显下降（n = 12：{pct(dis[(12, 0)])}；8 + 4：{pct(dis[(8, 4)])}；4 + 8：{pct(dis[(4, 8)])}）。"
        "如果只做'加干扰'这一个实验，会错误地得出'模型容易被无关信息干扰'；控制长度之后，真正的原因是"
        "**题面比训练时长**（长度泛化失败）。",
        f"演示 5：绕环本身没有问题（n = 4、k = 7：{pct(hops[(4, 7)]['acc'])}）。数字 8、9 从没在训练题里出现过，"
        f"模型遇到它们时最常写 {hops[(8, 8)]['len']} 步（占 {pct(hops[(8, 8)]['len_share'])}），写出的每一步都对"
        f"（{pct(hops[(8, 8)]['prefix_ok'])}），只是步数错了——失败的原因是'读不懂没见过的 token'，"
        "不是'不会走更多步'。如果不看推理链，你会把它误读成'推理链长度泛化失败'。",
        f"更隐蔽的是 n = 5、k = 9：答案准确率 {pct(hops[(5, 9)]['acc'])}，推理链全对 {pct(hops[(5, 9)]['chain_ok'])}——"
        "模型照样只走了 4 步，可 9 ≡ 4（mod 5），答案碰巧对了。只看最终答案的评测会把它算作'完全解决'。"
        "这是'答案对、推理错'的一个干净例子（讲义 §4.4）。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()