"""
=============================================================================
CS486 · 第 5 讲 · 代码 01：用强化学习修好"何时停下"——SFT vs RFT vs GRPO
=============================================================================

【这个脚本要回答的问题】（讲义 §4）
    第 4 讲发现：cot_L2 在 n = 12 的长题面上失败，主要原因是"走过头"——每一步查找都对，却没有在第 k 步停下。
    测试时计算（采样、投票、验证器）救不了这种系统性错误。现在我们**训练**模型来修好它，比较三种做法：

        SFT ：在 n = 4–12 的题上，用**标准答案**（正确的推理链）做监督微调——需要有人写好答案；
        RFT ：拒绝采样微调（rejection-sampling fine-tuning）：每道题让模型自己采样 8 次，
              **只拿自己做对的那些**做监督微调。STaR、ReST-EM 都是这个思路（讲义 §2）；
        GRPO：组相对策略优化（DeepSeekMath / DeepSeek-R1 用的 RL 算法）：每道题采样 8 次，
              比组内平均好的回答被鼓励、比平均差的被抑制——**做错的样本也提供信号**（讲义 §3）。

    三者都从第 1 讲的 cot_L2 出发、用同样的题目分布（n = 4–12，k = 1–7）、同样的步数（200）和学习率（1e-3），
    各跑 3 个随机种子。RFT 和 GRPO 只用"答案对不对"这一个信号（可验证奖励），从来不看标准推理链。

    演示 1：训练后在 n = 8（训练分布内）、12（RL 训练覆盖的最长题面）、14、16（更长，没训练过）上的贪心准确率。
    演示 2：训练过程——n = 12 的准确率怎样随步数变化；GRPO 的"零方差组"和熵。
    演示 3：失败类型——训练之后还剩下什么错误？

【运行】
    python3 01_rl_vs_sft.py           # 第一次约 1.5 分钟（9 个模型并行训练），之后约 15 秒
    python3 01_rl_vs_sft.py --quick   # 冒烟测试（几秒钟，数字没有意义）

【依赖】PyTorch + 本目录的 rl_lab.py（以及它导入的第 1、2、4 讲模块）。
"""

import argparse
import os
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 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

METHODS = ("sft", "rft", "grpo")
LABEL = {"base": "起点 cot_L2", "sft": "SFT（标准答案）", "rft": "RFT（自己做对的）", "grpo": "GRPO（组内比较）"}
NS = (8, 12, 14, 16)


def demo1(base, runs, per_k):
    """训练后的贪心准确率：每种方法 3 个种子的平均，以及每个种子的值。"""
    print_header("演示 1：训练后的贪心准确率（每个 n：k = 3–6 各 %d 道；RL 只训练过 n ≤ 12）" % per_k)
    sets = {n: rl.eval_set(n, per_k=per_k) for n in NS}
    acc = {"base": {n: [rl.greedy_accuracy(base, sets[n])] for n in NS}}
    for m in METHODS:
        acc[m] = {n: [rl.greedy_accuracy(model, sets[n]) for model, _ in runs[m]] for n in NS}
    rows = []
    for m in ("base",) + METHODS:
        cells = [LABEL[m]]
        for n in NS:
            v = acc[m][n]
            mean = sum(v) / len(v)
            cells.append(pct(mean) if len(v) == 1 else f"{pct(mean)}（{' / '.join(pct(x, 0) for x in v)}）")
        rows.append(cells)
    print_table(["模型"] + [f"n = {n}" for n in NS], rows, [18, 10, 24, 24, 24])
    print("  （括号里是 3 个种子各自的结果；第一列 n = 8 是第 1 讲的训练分布）")
    return acc


def demo2(runs):
    """训练曲线：n = 12 的准确率；GRPO 的训练统计。"""
    print_header("演示 2：训练过程（种子 0）——n = 12 测试集的贪心准确率随步数的变化")
    evals = {m: dict(runs[m][0][1]["evals"]) for m in METHODS}
    steps = sorted(evals["grpo"])
    shown = [s for s in steps if s in (0, 25, 50, 75, 100, 150, 200) or s == steps[-1]]
    rows = [[LABEL[m]] + [pct(evals[m][s]["acc"], 0) for s in shown] for m in METHODS]
    print_table(["方法 \\ 步数"] + [str(s) for s in shown], rows, [18] + [6] * len(shown))

    print("\n  GRPO（种子 0）每 25 步的训练统计（取这 25 步的平均）：")
    hist = runs["grpo"][0][1]["steps"]
    rows = []
    for lo in range(0, len(hist), 25):
        w = hist[lo:lo + 25]
        avg = {key: sum(h[key] for h in w) / len(w) for key in ("true", "zero", "ent", "len")}
        rows.append([f"{lo + 1}–{lo + len(w)}", pct(avg["true"]), pct(avg["zero"]), f"{avg['ent']:.4f}", f"{avg['len']:.2f}"])
    print_table(["步", "训练奖励", "零方差组", "每 token 熵", "回答长度"], rows, [10, 10, 10, 12, 10])
    rft = runs["rft"][0][1]["steps"]
    print(f"\n  对照：RFT（种子 0）前 25 步的训练奖励 {pct(sum(h['true'] for h in rft[:25]) / 25)}，"
          f"最后 25 步 {pct(sum(h['true'] for h in rft[-25:]) / 25)}")
    return evals, hist


def failure_counts(model, probs):
    """贪心解码，把每道题的回答归类（rl.classify）。"""
    res = tr.evaluate(model, probs, "cot")
    return Counter(rl.classify(p, rl.response_of(r["generated"])) for p, r in zip(probs, res))


def demo3(base, runs, per_k):
    print_header("演示 3：失败类型（种子 0，贪心解码）——训练之后还剩下什么错误？")
    out = {}
    for n in (12, 16):
        probs = rl.eval_set(n, per_k=per_k)
        rows = []
        for m in ("base",) + METHODS:
            model = base if m == "base" else runs[m][0][0]
            c = failure_counts(model, probs)
            out[(n, m)] = c
            rows.append([LABEL[m]] + [c.get(t, 0) for t in rl.FAILURE_TYPES])
        print(f"\n  n = {n}（{len(probs)} 道题）：")
        print_table(["模型", "对", "走过头", "停早了", "查找错", "抄错", "没写 |"], rows, [18, 6, 8, 8, 8, 6, 8])
    print("  （没写 |：写到长度上限 k + 5 还没停下，也就是走过头至少 5 步；走过头：写了 | 但推理链超过 k 步）")
    return out


def main():
    ap = argparse.ArgumentParser(description="SFT vs RFT vs GRPO")
    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；三种方法 × 3 个种子")
    base = tr.get_model("cot_L2", quick=args.quick, verbose=not args.quick)[0]
    jobs = rl.main_jobs()
    if args.quick:
        jobs = [j for j in jobs if j["seed"] == 0]
    results = rl.load_runs(jobs, quick=args.quick)
    runs = {m: [res for j, res in zip(jobs, results) if j["method"] == m] for m in METHODS}

    acc = demo1(base, runs, per_k)
    evals, hist = demo2(runs)
    fails = demo3(base, runs, per_k)

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    mean = {m: {n: sum(acc[m][n]) / len(acc[m][n]) for n in NS} for m in acc}
    b12 = fails[(12, "base")]
    g16 = fails[(16, "grpo")]
    first = sum(h["zero"] for h in hist[:25]) / 25
    last = sum(h["zero"] for h in hist[-25:]) / 25
    rft16 = " / ".join(pct(x, 0) for x in acc["rft"][16])
    print_observations([
        f"三种方法都修好了 n = 12（演示 1）：起点 {pct(mean['base'][12])}，SFT / RFT / GRPO 训练后（3 个种子的平均）"
        f"{pct(mean['sft'][12])} / {pct(mean['rft'][12])} / {pct(mean['grpo'][12])}，而且只用了 25–50 步（演示 2）。"
        f"起点在 n = 12 上的 {200 - b12['correct']} 个错误里，{b12['malformed']} 个是一直没写 '|'、{b12['overshoot']} 个是走过头、"
        f"{b12['lookup']} 个是查找错（演示 3）；训练后全部消失。RFT 和 GRPO 从来没看过标准推理链，"
        "只知道'答案对不对'，就学会了'在第 k 步停下'——这正是第 4 讲里测试时计算（投票、验证器）救不了的系统性错误。",
        f"在没训练过的更长题面上（n = 14 / 16），三种方法都有提高、但远没有解决：SFT {pct(mean['sft'][14])} / {pct(mean['sft'][16])}，"
        f"GRPO {pct(mean['grpo'][14])} / {pct(mean['grpo'][16])}，RFT {pct(mean['rft'][14])} / {pct(mean['rft'][16])}。"
        f"种子之间差别很大（RFT 在 n = 16 上三个种子是 {rft16}），3 个种子下 SFT 和 GRPO 的差别在噪声范围内。"
        "'SFT 只会记忆、RL 才会泛化'在这里**没有**出现：我们的标准答案完美、和模型该写的完全一致，SFT 没有吃亏的理由；"
        "那类结论针对的是 SFT 数据和模型自己会写的东西不一样的情形（讲义 §4.4）。",
        f"GRPO 的学习信号非常稀疏（演示 2）：前 25 步平均有 {pct(first)} 的题组'8 个样本全对或全错'，优势全为 0、对梯度毫无贡献；"
        f"最后 25 步是 {pct(last)}。起点模型在温度 1.0 下几乎是确定性的（前 25 步回答 token 的平均熵 "
        f"{sum(h['ent'] for h in hist[:25]) / 25:.4f} nats），"
        "大多数题要么总是对、要么总是错——**只有'有时对'的题才提供信号**。这就是 DAPO 的'动态采样'（过滤掉零方差组、补采新题）"
        "和'按难度挑题'的动机。熵在训练中继续下降，代码 02 看这对 pass@k 意味着什么。",
        f"RFT 只学自己做对的样本（不惩罚做错的），在 n = 12 上和 GRPO 一样好，在更长的题面上却更差、更不稳定。"
        "一个可能的原因：GRPO 会把'没停下'的回答的概率**压低**，RFT 只把正确回答的概率抬高；"
        "在训练分布内两者等价（正确回答的概率抬到接近 1，其他回答自然被挤掉），在分布外就不一样了。这里没有单独检验（练习 2）。",
        f"n = 16 上剩下的错误仍然主要是'停不下来'：GRPO 的 {200 - g16['correct']} 个错误里有 {g16['malformed']} 个没写 '|'、"
        f"{g16['overshoot']} 个走过头，查找错只有 {g16['lookup']} 个。在 n ≤ 12 上学到的'何时停下'只有一部分迁移到了更长的题面——"
        "长度泛化的瓶颈是停止，不是查找。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()