"""
=============================================================================
CS486 · 第 3 讲 · 代码 02：思维链的"步长"——每一步能跨多远？
=============================================================================

【这个脚本要回答的问题】（讲义 §5）
    思维链把一个 k 跳的问题拆成 k 个 1 跳的小步骤，每一步都在一次前向的能力之内（代码 01：2 层 = 绑定 + 查找 = 1 跳）。
    如果推理链写得"更省"——每一步跨 s 跳（只写第 s、2s、3s… 跳到达的字母）——会怎样？
        s = 1：普通思维链（第 1 讲的 cot_L2，100%）
        s = 2：每一步要在一次前向里做 2 跳——2 层模型做不到（第 1 讲代码 02），4 层模型表达得了但学得很慢
        s = 3：每一步 3 跳
    真实 LLM 里的对应问题：推理步骤写得太"跳"（省略中间步骤）、蒸馏时压缩推理链、
    "学会跳步"（Liu et al. 2024）等——每一步的粒度要和模型一次前向的能力匹配。

    演示 1：步长 s = 1、2、3 的 2 层模型，以及 s = 2 的 4 层模型，在 k = 1–7 上的准确率。
    演示 2：推理链第一步的准确率——k = 1 时第一步只要 1 跳，模型学会了吗？

【运行】
    python3 02_step_size.py           # 第一次约 1.5–2 分钟（并行训练 3 个模型；s = 1 用第 1 讲的缓存），之后约 5 秒
    python3 02_step_size.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

JOBS = [
    {"kind": "stride", "stride": 2, "layers": 2, "steps": 2000, "seed": 0},
    {"kind": "stride", "stride": 3, "layers": 2, "steps": 2000, "seed": 0},
    {"kind": "stride", "stride": 2, "layers": 4, "steps": 3000, "seed": 0},
]


def evaluate_stride(model, s, k, count):
    """
    在 n = 8、k 跳的题上评测步长 s 的思维链模型。

    【返回】(答案准确率, 推理链第一个字母的准确率)
    """
    probs = hw.make_eval_set(8, k, count, seed=5)
    need = len(cl.stride_trace(probs[0], s)) + 3
    outs = tr.generate(model, [p.prompt for p in probs], need)
    acc = sum(hw.parse_answer(o, "cot") == p.answer for o, p in zip(outs, probs)) / count
    first = sum(o[:1] == cl.stride_trace(p, s)[0] for o, p in zip(outs, probs)) / count
    return acc, first


def main():
    ap = argparse.ArgumentParser(description="思维链的步长")
    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("准备：同样的训练题，只改推理链的步长")
    for s in (1, 2, 3):
        p = hw.make_eval_set(8, 5, 1, seed=5)[0]
        print(f"  s = {s}：{p.prompt} → {cl.stride_trace(p, s)}|{p.answer}$")
    base = tr.get_model("cot_L2", quick=args.quick)[0]
    models = cl.load_jobs(JOBS, quick=args.quick)
    entries = [("s = 1，2 层（cot_L2）", base, 1)] + \
              [(f"s = {j['stride']}，{j['layers']} 层，{j['steps']} 步", m, j["stride"]) for j, m in zip(JOBS, models)]

    print_header(f"演示 1：答案准确率（n = 8，每格 {count} 道题）")
    res = {}
    for label, model, s in entries:
        res[label] = [evaluate_stride(model, s, k, count) for k in range(1, 8)]
    rows = [[f"k = {k}"] + [pct(res[label][k - 1][0]) for label, _, _ in entries] for k in range(1, 8)]
    print_table(["跳数"] + [label for label, _, _ in entries], rows, [8, 22, 22, 22, 22])

    print_header("演示 2：推理链第一个字母的准确率（k = 1 时第一步只需要 1 跳）")
    rows = []
    for k in (1, 2, 3, 4):
        row = [f"k = {k}"]
        for label, _, s in entries:
            row.append(f"{pct(res[label][k - 1][1])}（{min(s, k)} 跳）")
        rows.append(row)
    print_table(["跳数"] + [label for label, _, _ in entries], rows, [8, 22, 22, 22, 22])
    print("（括号里是第一步需要跨的跳数：min(s, k)。）")

    if args.quick:
        print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    lab = [label for label, _, _ in entries]
    avg = {label: sum(a for a, _ in res[label]) / 7 for label in lab}
    print_observations([
        f"步长 1 的思维链 {pct(avg[lab[0]])}；步长 2、3 的 2 层模型平均只有 {pct(avg[lab[1]])} 和 {pct(avg[lab[2]])}，"
        "在随机水平附近。思维链本身不是魔法：每一步必须在一次前向的能力之内（2 层 = 1 跳）。",
        f"步长 2 的 4 层模型（3000 步）平均 {pct(avg[lab[3]])}：4 层**表达得了**每步 2 跳（第 1 讲代码 02 的 --deep），"
        "但在这个训练预算内没学会——和第 1 讲'4 层模型学 2 跳很慢'一致。",
        f"最能说明问题的是演示 2 的 k = 1 一行：步长 2、3 时，k = 1 的推理链第一步只需要 1 跳，可 2 层模型的准确率只有 "
        f"{pct(res[lab[1]][0][1])} / {pct(res[lab[2]][0][1])}——连 1 跳都没学会。这和第 1 讲'混合训练的直接回答模型连 1 跳都学不会'"
        "是同一个现象：训练数据里大部分步骤学不会，淹没了能学会的那部分。步长太大不只是让大步骤失败，而是让整个模型学不会。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()