"""
=============================================================================
CS486 · 第 1 讲 · 代码 02：一次前向计算能做几跳？——表达能力（expressivity）与可学习性（learnability）
=============================================================================

【这个脚本要回答的问题】（讲义 §3.5）
    代码 01 里，直接回答的模型连 k = 1 都没学会——可是 k = 1 和思维链的第一步是完全相同的计算，
    思维链模型显然会做。这说明"做不到"可能有两种完全不同的原因：
        (a) **表达能力**：这个结构的模型，无论参数取什么值，都算不出这个函数；
        (b) **可学习性**：存在能算出它的参数，但梯度下降在给定的数据和步数里找不到。
    研究里把这两者混为一谈是常见的错误。这个脚本把它们拆开：

    演示 1：只用 k = 1 的题训练直接回答模型（2 层）；再只用 k = 2 的题训练一个（2 层）。
            → 1 跳学得会吗？2 跳呢？学会 1 跳的模型遇到 k = 2、3 的题会怎样？
    演示 2：同样的"1 跳查找"在三种训练条件下的表现：单独训练、混在多跳题里训练、作为思维链的一步。
            → 为什么混合训练连 1 跳都学不会？（可学习性）
    演示 3（--deep，可选，约 6 分钟）：训练 10000 步，比较 2 层和 4 层模型能不能学会 2 跳。
            → 更深的模型能在一次前向里串起更多跳吗？代价是什么？

【运行】
    python3 02_hops_per_forward_pass.py            # 演示 1–2，第一次约 1 分钟（训练 2 个模型），之后几秒
    python3 02_hops_per_forward_pass.py --deep     # 另加演示 3（再训练 2 个模型，约 6 分钟）
    python3 02_hops_per_forward_pass.py --quick    # 冒烟测试（加 --deep 也可以），几秒钟，数字没有意义
    演示 2 用到代码 01 的 direct_L2 和 cot_L2；如果还没有缓存，会先训练它们（约 1 分钟）。

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

import argparse
import time

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


def hop_profile(model, mode, ks, count, n=8):
    """
    对每个 k，统计模型的答案落在"第几跳"上。

    【思路】一个只会走 1 跳的模型，在 k = 3 的题上的答案应该恰好是 walk[1]。
        统计"答案 = 第 j 跳的位置"的比例（j = 1、2、3），能直接看出模型实际执行的跳数，
        而不是只看"对 / 错"。
    【参数】model、mode；ks：要测的跳数列表；count：每个 k 的题数；n：环大小
    【返回】dict：k → {"acc": 准确率, "hop1": 答案 = 第 1 跳位置的比例, "hop2": …, "hop3": …}
    【边界】n = 8 时第 j 跳的位置互不相同（j < 8），所以这些比例不会重复计数。
    """
    out = {}
    for k in ks:
        probs = hw.make_eval_set(n, k, count, seed=2)
        res = tr.evaluate(model, probs, mode)
        d = {"acc": tr.accuracy(res)}
        for j in (1, 2, 3):
            d[f"hop{j}"] = sum(r["pred"] == _walk_far(p, j) for p, r in zip(probs, res)) / len(probs)
        out[k] = d
    return out


def _walk_far(prob, j):
    """
    从起点走 j 跳到达的字母（j 可以大于题目的 k）。

    【思路】prob.walk 只记录到第 k 跳；要知道"如果模型走了 j 跳会停在哪"，就从题面重新求解 j 跳。
    """
    succ = {st[0]: st[2] for st in prob.statements}
    x = prob.start
    for _ in range(j):
        x = succ[x]
    return x


def hop_index_distribution(model, mode, n, k, count, seed=21):
    """
    模型的答案是"从起点走几跳"的位置？返回 j = 0 … n-1 每一类的比例。

    【思路】在 n 个节点的环上，每个字母都恰好是"从起点走 j 跳"的位置（j = 0 是起点本身，
        j = n-1 是起点的前驱）。把答案换算成 j，就能看出模型在"怎样猜"：
        比如从不回答 j = 0、1、n-1，说明它用一次查找就能确定的信息**排除**了这几个字母。
    【返回】list，第 j 个元素是答案恰好为第 j 跳位置的比例（解析失败的不计入，所以总和可能略小于 1）
    """
    probs = hw.make_eval_set(n, k, count, seed=seed)
    res = tr.evaluate(model, probs, mode)
    counts = [0] * n
    for p, r in zip(probs, res):
        for j in range(n):
            if r["pred"] == _walk_far(p, j):
                counts[j] += 1
                break
    return [c / len(probs) for c in counts]


def demo1_fixed_k(quick, count):
    print_header("演示 1：只用一种跳数训练的直接回答模型（2 层）——一次前向能做几跳？")
    rows, profiles = [], {}
    for name, label in (("direct_k1_L2", "只训练 k = 1"), ("direct_k2_L2", "只训练 k = 2")):
        model, info, tcfg = tr.get_model(name, quick=quick)
        prof = hop_profile(model, "direct", (1, 2, 3), count)
        profiles[name] = prof
        for k in (1, 2, 3):
            d = prof[k]
            rows.append([label, f"k = {k}", pct(d["acc"]), pct(d["hop1"]), pct(d["hop2"]), pct(d["hop3"])])
    print()
    print_table(["模型", "测试跳数", "准确率", "答案=第1跳", "答案=第2跳", "答案=第3跳"], rows,
                [14, 10, 8, 11, 11, 11], ["<", "<", ">", ">", ">", ">"])
    print(f"（n = 8，每格 {count} 道题；随机猜一个字母约 12.5%。'答案=第 j 跳'：模型的答案恰好是从起点走 j 跳的位置。）")
    return profiles


def demo2_learnability(quick, count, k1_profile):
    print_header("演示 2：同一个'1 跳查找'，在三种训练条件下")
    direct_mixed, _, _ = tr.get_model("direct_L2", quick=quick)
    cot, _, _ = tr.get_model("cot_L2", quick=quick)
    probs = hw.make_eval_set(8, 1, count, seed=2)
    acc_mixed = tr.accuracy(tr.evaluate(direct_mixed, probs, "direct"))

    # 思维链模型的"第 1 个推理字母"就是 1 跳查找；对 k = 1–7 的题都统计它
    first_ok = total = 0
    for k in range(1, 8):
        ps = hw.make_eval_set(8, k, count // 3, seed=2)
        for p, r in zip(ps, tr.evaluate(cot, ps, "cot")):
            total += 1
            first_ok += r["generated"][:1] == p.walk[1]
    # 思维链的"每一步"都是 1 跳查找：统计推理链里每个字母都对的比例（逐步）
    step_ok = step_total = 0
    for k in range(1, 8):
        ps = hw.make_eval_set(8, k, count // 3, seed=3)
        for p, r in zip(ps, tr.evaluate(cot, ps, "cot")):
            for i, ch in enumerate(p.trace):
                step_total += 1
                step_ok += r["generated"][i:i + 1] == ch
    rows = [
        ["单独训练（只有 k = 1 的题）", "direct_k1_L2", pct(k1_profile[1]["acc"])],
        ["混在多跳题里训练（k = 1 … n-1）", "direct_L2", pct(acc_mixed)],
        ["作为思维链的第 1 步", "cot_L2", pct(first_ok / total)],
        ["作为思维链的每一步（逐个字母）", "cot_L2", pct(step_ok / step_total)],
    ]
    print_table(["训练条件", "模型", "1 跳查找的准确率"], rows, [34, 14, 16], ["<", "<", ">"])
    return acc_mixed, first_ok / total, step_ok / step_total


def demo3_deep(quick, count):
    print_header("演示 3（--deep）：训练 10000 步（学习率 1e-3），2 层 vs 4 层，只用 k = 2 的题")
    curves, finals = {}, {}
    for name in ("direct_k2_L2_long", "direct_k2_L4_long"):
        model, info, tcfg = tr.get_model(name, quick=quick)
        curves[name] = {h["step"]: h["probe_acc"] for h in info["history"] if h["probe_acc"] is not None}
        finals[name] = hop_profile(model, "direct", (1, 2, 3), count)
    steps = sorted(curves["direct_k2_L4_long"])
    print("\n学习曲线（探针：n = 8、k = 2 的 140 道题）：")
    rows = [[s, pct(curves["direct_k2_L2_long"].get(s)), pct(curves["direct_k2_L4_long"].get(s))] for s in steps]
    print_table(["训练步数", "2 层", "4 层"], rows, [10, 10, 10])
    print(f"\n训练结束后（n = 8，每格 {count} 道题）：")
    rows = []
    for name, label in (("direct_k2_L2_long", "2 层"), ("direct_k2_L4_long", "4 层")):
        for k in (1, 2, 3):
            d = finals[name][k]
            rows.append([label, f"k = {k}", pct(d["acc"]), pct(d["hop1"]), pct(d["hop2"]), pct(d["hop3"])])
    print_table(["模型", "测试跳数", "准确率", "答案=第1跳", "答案=第2跳", "答案=第3跳"], rows,
                [8, 10, 8, 11, 11, 11], ["<", "<", ">", ">", ">", ">"])
    print(f"\n答案落在'第 j 跳'位置的比例（n = 8、k = 2，{count} 道题；j = 0 是起点，j = 7 是起点的前驱）：")
    dists = {}
    rows = []
    for name, label in (("direct_k2_L2", "2 层，2000 步"), ("direct_k2_L2_long", "2 层，10000 步"),
                        ("direct_k2_L4_long", "4 层，10000 步")):
        model = tr.get_model(name, quick=quick, verbose=False)[0]
        dists[name] = hop_index_distribution(model, "direct", 8, 2, count)
        rows.append([label] + [pct(x, 0) for x in dists[name]])
    print_table(["模型"] + [f"j={j}" for j in range(8)], rows, [16] + [6] * 8)
    return curves, finals, dists


def main():
    ap = argparse.ArgumentParser(description="一次前向计算能做几跳？")
    ap.add_argument("--quick", action="store_true", help="冒烟测试：每个模型只训练 60 步")
    ap.add_argument("--deep", action="store_true", help="另外运行演示 3（训练两个 10000 步的模型，约 6 分钟）")
    args = ap.parse_args()
    t0 = time.perf_counter()
    count = 30 if args.quick else 300

    profiles = demo1_fixed_k(args.quick, count)
    acc_mixed, first_acc, step_acc = demo2_learnability(args.quick, count, profiles["direct_k1_L2"])
    if args.deep:
        curves, finals, dists = demo3_deep(args.quick, count)
    else:
        print("\n（加 --deep 运行演示 3：2 层 vs 4 层，训练 10000 步，约 6 分钟。）")

    if args.quick:
        print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return

    p1, p2 = profiles["direct_k1_L2"], profiles["direct_k2_L2"]
    obs = [
        f"只用 k = 1 训练的 2 层模型学会了 1 跳（准确率 {pct(p1[1]['acc'])}）；把它用在 k = 2、3 的题上，"
        f"它的答案有 {pct(p1[2]['hop1'])}、{pct(p1[3]['hop1'])} 恰好是'第 1 跳'的位置——"
        "它学会的是'走一步'，完全忽略了题面里的跳数。",
        f"只用 k = 2 训练的 2 层模型没有学会 2 跳（准确率 {pct(p2[2]['acc'])}，在随机猜的 1/8–1/7 之间）。"
        "讲义 §3.5：最直接的查找机制每一跳需要先'把陈述的两端绑在一起'、再'按名字查找'，"
        "2 层只够做 1 跳（--deep 的长训练会看到，2 层模型能靠'排除法'超过随机，但学不会 2 跳）。",
        f"演示 2：同一个 1 跳查找，单独训练时 {pct(p1[1]['acc'])}，作为思维链的一步时 {pct(step_acc)}，"
        f"但混在多跳题里训练直接回答时只有 {pct(acc_mixed)}。2 层模型**表达得了** 1 跳，"
        "混合训练却没学会——这是可学习性问题：k ≥ 2 的题（训练分布里约占 78%）对 2 层模型是学不会的噪声，"
        "梯度信号被淹没；而思维链把每道题拆成 k 个 1 跳，每个 token 都是一个可学的小任务（稠密的中间监督）。",
    ]
    if args.deep:
        c2, c4 = curves["direct_k2_L2_long"], curves["direct_k2_L4_long"]
        f2, f4 = finals["direct_k2_L2_long"], finals["direct_k2_L4_long"]
        last = max(c4)
        obs.append(
            f"演示 3：训练 10000 步后，4 层模型在 k = 2 上达到 {pct(f4[2]['acc'])}，2 层模型是 {pct(f2[2]['acc'])}。"
            f"学习曲线显示 2 层模型从第 6000 步起停在 30% 左右的平台（第 {last} 步 {pct(c2[last])}），"
            f"4 层模型则一直在上升、训练结束时还没学完（第 {last} 步 {pct(c4[last])}）。"
            "更深的模型能在一次前向里串起更多跳，但学得很慢：思维链模型 750 步就学会了 1–7 跳（代码 01 演示 3），"
            f"4 层模型 10000 步还没学完 2 跳。而且它学的是'恰好 2 跳'：k = 1、3 的题上准确率是 {pct(f4[1]['acc'])}、{pct(f4[3]['acc'])}。")
        d2 = dists["direct_k2_L2_long"]
        obs.append(
            f"2 层模型训练 10000 步后的 {pct(finals['direct_k2_L2_long'][2]['acc'])} 并不是'学会了一部分 2 跳'：答案分布显示它"
            f"几乎从不回答起点（j = 0：{pct(d2[0])}）、起点的后继（j = 1：{pct(d2[1])}）和起点的前驱（j = 7：{pct(d2[7])}）——"
            "这三个字母恰好是**一次查找**就能确定的；它把它们排除掉，再在剩下的 5 个字母里带偏好地猜。"
            "高于随机的准确率可以来自'排除法'，不一定来自目标机制——看准确率之外的分布才能分辨。")
    print_observations(obs)
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()