"""
=============================================================================
第 6 讲 · 02：策略梯度 —— REINFORCE、基线、RLOO 与 GRPO
=============================================================================

对应讲义：lectures/lecture06_scaling_rl.md 第 3、4 节（DeepSeekMath 的 GRPO）

-----------------------------------------------------------------------------
【共同的设定：RL with verifiable rewards（RLVR）】
-----------------------------------------------------------------------------
奖励 r = 1（答案正确）或 0（错误），由程序核对，不需要 reward model。
每个训练步：随机抽 16 道训练题，每道题采样 G = 8 个解答（一共 128 个），按下面的公式算梯度，更新一次参数：

        梯度 = (1/128) Σ_i  A_i × Σ_t ∇ log π(a_t | 情境_t)

    其中 a_t 是第 i 个解答在第 t 步选的做法，A_i 是这个解答的"优势"（advantage）——它比"平常"好多少。
    不同的方法，区别只在 A_i 怎么算：
        · 只在答对的样本上 SFT ：A_i = r_i（不减任何东西；等价于 STaR 的在线版本 / rejection sampling fine-tuning）
        · REINFORCE + 全局基线 ：A_i = r_i − b，b 是过去奖励的滑动平均
        · RLOO                  ：A_i = r_i − （同一道题其他 G − 1 个样本的平均奖励）
        · GRPO                  ：A_i = (r_i − 组内平均) / (组内标准差 + 1e−6)
    同一道题的一组样本互相比较 —— 这就是"group relative"的意思。GRPO 不需要训练一个价值网络（critic）来估计基线。

运行时间：约 1 秒（只用标准库）。
"""

import math
import random

from trainable_llm import (TrainableLLM, add_grad, apply_grads, context_report, entropy_of, evaluate,
                           make_problem_set, mean, pad, pass_at_k, probs_for, step_records)

GROUP = 8
BATCH = 16
STEPS = 60
LR = 2.0


def advantages(method, rewards, state):
    """
    按方法计算一组样本的优势。

    【参数】method：方法名；rewards：同一道题的 G 个奖励；state：dict，保存全局基线等跨步的状态
    【返回】list of float（和 rewards 等长）
    【边界】GRPO 里一组奖励全相同时（全对或全错），标准差为 0，优势全为 0 —— 这组样本对梯度**没有贡献**。
        第 03 号文件的 dynamic sampling 就是为了处理这种"浪费掉的"样本。
    """
    g = len(rewards)
    if method == "sft_correct":
        return list(rewards)
    if method == "reinforce_baseline":
        b = state.get("baseline", mean(rewards))
        return [r - b for r in rewards]
    if method == "rloo":
        total = sum(rewards)
        return [r - (total - r) / (g - 1) for r in rewards]
    if method == "grpo":
        mu = mean(rewards)
        sd = math.sqrt(mean((r - mu) ** 2 for r in rewards))
        return [(r - mu) / (sd + 1e-6) for r in rewards]
    raise ValueError(method)


def train(method, train_set, test_set, seed, lr=LR):
    """
    用某种优势的定义训练 STEPS 步，返回 (模型, 学习曲线)。

    【返回】curve：list of (训练步, 测试集 pass@1)
    """
    rng = random.Random(seed)
    model = TrainableLLM(1e9)
    state, curve = {}, [(0, evaluate(model, test_set, rng)[1])]
    for step in range(1, STEPS + 1):
        batch = rng.sample(train_set, BATCH)
        grads, all_rewards = {}, []
        n_samples = BATCH * GROUP
        for p in batch:
            samples = [model.generate(p, rng) for _ in range(GROUP)]
            rewards = [1.0 if s.correct else 0.0 for s in samples]
            all_rewards.extend(rewards)
            for s, a in zip(samples, advantages(method, rewards, state)):
                if a == 0.0:
                    continue
                for (key, action), probs in zip(step_records(p, s), probs_for(model, p, s)):
                    add_grad(grads, key, action, probs, a / n_samples)
        apply_grads(model, grads, lr)
        state["baseline"] = 0.9 * state.get("baseline", mean(all_rewards)) + 0.1 * mean(all_rewards)
        if step % 15 == 0:
            curve.append((step, evaluate(model, test_set, rng)[1]))
    return model, curve


def pass_k_at_temperature(model, problems, rng, temperature, n=16, ks=(1, 16)):
    """在指定温度下采样 n 次，返回 {k: pass@k}。"""
    counts = [sum(model.generate(p, rng, temperature).correct for _ in range(n)) for p in problems]
    return {k: mean(pass_at_k(n, c, k) for c in counts) for k in ks}


def main():
    train_set = make_problem_set(random.Random(150), n_problems=300, min_steps=3, max_steps=8)
    test_set = make_problem_set(random.Random(151), n_problems=200, min_steps=3, max_steps=8)
    methods = (("只在答对的样本上 SFT", "sft_correct"), ("REINFORCE + 全局基线", "reinforce_baseline"),
               ("RLOO", "rloo"), ("GRPO", "grpo"))
    results = {key: train(key, train_set, test_set, seed=152) for _, key in methods}
    results["rloo_x4"] = train("rloo", train_set, test_set, seed=152, lr=4 * LR)

    print("=" * 78)
    print(f"演示 1：四种优势的定义 —— 测试集 pass@1 随训练步数的变化（每步 {BATCH} 题 × {GROUP} 个样本，学习率 {LR}）")
    print("=" * 78)
    steps = [s for s, _ in results["grpo"][1]]
    print(pad("方法", 22) + "".join(pad(f"第 {s} 步", 10, "right") for s in steps))
    for name, key in methods + (("RLOO（学习率 × 4）", "rloo_x4"),):
        print(pad(name, 22) + "".join(pad(f"{acc:.3f}", 10, "right") for _, acc in results[key][1]))
    # 训练开始时，各种方法的优势平均有多大？
    rng = random.Random(154)
    base = TrainableLLM(1e9)
    groups = [[1.0 if base.generate(p, rng).correct else 0.0 for _ in range(GROUP)] for p in train_set[:100]]
    mags = {key: mean(abs(a) for g in groups for a in advantages(key, g, {"baseline": mean(sum(groups, []))}))
            for _, key in methods}
    print("\n训练开始时，优势的平均绝对值 |A|：" + "，".join(f"{name} {mags[key]:.3f}" for name, key in methods))
    print("\n观察：")
    print("  · 四种方法都只用\"答案对不对\"作为奖励，都能把 pass@1 从 0.12 提高到 0.8 以上：RLVR 不需要 reward model。")
    print("  · 同样的学习率下，GRPO 学得最快。但这主要是因为它**除以了组内的标准差**：优势的平均幅度是 RLOO 的两倍，")
    print("    而且在\"8 个里只对 1 个\"这样的组里，那个答对的样本会得到很大的优势（标准差很小）。")
    print("    把 RLOO 的学习率乘以 4，它就和 GRPO 一样快甚至更快 —— 归一化本质上是一种按题目自适应的步长。")
    print("    后来的分析（Dr. GRPO 等）也指出，这个归一化会让模型对\"太难\"和\"太简单\"的题过度敏感，带来偏差。")
    print("  · \"只在答对的样本上 SFT\"没有减去基线：答错的样本被直接忽略，不会被\"推开\"。")
    print("    减去基线之后（REINFORCE、RLOO、GRPO），答错的样本得到负的优势；在\"大家都能答对\"的简单题上，")
    print("    答对的样本几乎不再被鼓励 —— 梯度集中在\"有的对、有的错\"的题上。\n")

    # 演示 2：RL 学到了什么？和"降低温度"比较
    print("=" * 78)
    print("演示 2：RL 学到了什么？—— 和\"直接把原模型的温度调低\"比较")
    print("=" * 78)
    base = TrainableLLM(1e9)
    grpo_model = results["grpo"][0]
    rng = random.Random(153)
    rows = [("原模型，温度 1.0", base, 1.0), ("原模型，温度 0.7", base, 0.7), ("原模型，温度 0.4", base, 0.4),
            ("GRPO 训练后，温度 1.0", grpo_model, 1.0)]
    print(pad("模型", 24) + pad("pass@1", 10, "right") + pad("pass@16", 10, "right") + pad("平均熵", 10, "right"))
    for name, m, t in rows:
        r = pass_k_at_temperature(m, test_set, rng, t)
        print(pad(name, 24) + pad(f"{r[1]:.3f}", 10, "right") + pad(f"{r[16]:.3f}", 10, "right")
              + pad(f"{entropy_of(m, test_set[:60], rng):.3f}" if t == 1.0 else "—", 10, "right"))
    print("\n几个情境下选 correct 的概率：")
    print(pad("情境", 26) + pad("原模型", 10, "right") + pad("GRPO 训练后", 14, "right"))
    for (ctx, p0), p1 in zip(context_report(base).items(), context_report(grpo_model).values()):
        print(pad(ctx, 26) + pad(f"{p0:.3f}", 10, "right") + pad(f"{p1:.3f}", 14, "right"))
    print("\n观察：")
    print("  · 降低温度能提高一点 pass@1（模型\"更确定\"了），降到 0.4 时 pass@16 开始下降（多样性没了），")
    print("    而且救不了模型\"最相信错误做法\"的情境：乘两位数时 bug 的 logit 比 correct 高，降温只会更常犯错。")
    print("  · GRPO 训练后，每个情境都几乎总是选 correct，pass@1 接近 1：它不只是\"变得更确定\"，")
    print("    还把原本 bug 占优的情境改成了 correct 占优 —— 奖励告诉了它哪种做法才是对的。")
    print("  · 熵从 0.84 降到 0.04：策略几乎变成了确定性的（entropy collapse）。在这个 toy world 里这不是坏事，")
    print("    因为每一步只有一种正确的做法；但它意味着模型不再探索（第 03 号文件）。")
    print("  · 在这里 pass@16 也从 0.64 升到了 1.0 —— 这和真实 LLM 上的观察**不同**。DeepSeekMath 发现 RL 显著提高了")
    print("    Maj@K，却没有明显提高 Pass@K；Yue et al. 2025 甚至发现 k 很大时原模型的 pass@k 更高，")
    print("    认为 RL 主要是\"把偶尔能做对变成经常做对\"，而不是扩展能力的边界。")
    print("    toy world 里每一步只有一种正确做法、参数在所有题之间共享，学会了就处处适用，所以看不到这种现象。\n")


if __name__ == "__main__":
    main()