"""
=============================================================================
第 6 讲 · 03：DAPO —— 大规模 RL 训练的几个技巧
=============================================================================

对应讲义：lectures/lecture06_scaling_rl.md 第 5 节（Yu et al. 2025，DAPO）

-----------------------------------------------------------------------------
【从 GRPO 出发】
-----------------------------------------------------------------------------
真实的 GRPO（和 PPO）每采一批数据，会用它**更新好几次**参数（这里 μ = 4 次）。更新几次之后，当前策略 π_θ
已经和采样时的旧策略 π_old 不同了，所以每一步的梯度要乘上重要性比值 ratio = π_θ(a) / π_old(a)，
并且**裁剪**（clip）：ratio 超出 [1 − ε_low, 1 + ε_high] 就不再往那个方向推，防止一批数据把策略推得太远。

        A > 0 且 ratio > 1 + ε_high  →  这一步不再贡献梯度（已经鼓励够了）
        A < 0 且 ratio < 1 − ε_low   →  这一步不再贡献梯度（已经打压够了）
        其他情况                      →  梯度 = A × ratio × ∇ log π_θ(a)

DAPO 在此基础上提出了四个技巧，这里实现前三个：
    1. clip-higher        ：ε_low = 0.2 不变，ε_high 从 0.2 放宽到 0.28。
                            低概率的动作（例如难情境里的 correct）每批能被提高得更多，探索保持得更久。
    2. dynamic sampling   ：一组 G 个样本如果全对或全错，优势全为 0，对梯度没有贡献。
                            DAPO 不断采样、丢掉这样的组，直到凑满一批"有对有错"的组。
    3. token-level loss   ：GRPO 先在每个回答内部对 token 求平均，再对回答求平均 —— 长回答里的每个 token 权重更小。
                            DAPO 直接对所有 token 求平均，长回答和短回答里的 token 一视同仁。
    4. overlong reward shaping（不实现）：对超长被截断的回答，不直接给 −1，而是在一个缓冲区间里线性地惩罚长度。
                            toy world 的解答不会"写太长被截断"，所以无法演示。
    DAPO 还去掉了 GRPO 里的 KL 惩罚；这里的 GRPO 基线本来就没有加 KL。

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

import math
import random

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

GROUP = 8
BATCH = 16          # 每批多少组
MU = 4              # 每批数据更新几次
LR = 4.0
STEPS = 30


def group_advantages(rewards):
    """GRPO 的组内相对优势：(r − 均值) / (标准差 + 1e−6)。全对或全错时全为 0。"""
    mu = mean(rewards)
    sd = math.sqrt(mean((r - mu) ** 2 for r in rewards))
    return [(r - mu) / (sd + 1e-6) for r in rewards]


def collect_batch(model, train_set, rng, dynamic):
    """
    采一批数据。

    【返回】(groups, n_generated, n_zero_groups)
        groups：list of (problem, samples, advantages)；n_generated：一共采了多少个样本（成本）；
        n_zero_groups：采到的组里，有多少组优势全为 0（全对或全错）
    【思路】dynamic=False：随机抽 BATCH 道题，每题采 G 个样本，不管有没有信息量都放进批里。
        dynamic=True：不断随机抽题、采样，丢掉优势全为 0 的组，直到凑满 BATCH 个"有对有错"的组
        （最多尝试 4 × BATCH 道题，防止模型太强 / 太弱时无限循环）。
    """
    groups, n_generated, n_zero, tries = [], 0, 0, 0
    while len(groups) < BATCH and tries < 4 * BATCH:
        tries += 1
        p = rng.choice(train_set)
        samples = [model.generate(p, rng) for _ in range(GROUP)]
        n_generated += GROUP
        rewards = [1.0 if s.correct else 0.0 for s in samples]
        adv = group_advantages(rewards)
        zero = all(a == 0.0 for a in adv)
        n_zero += zero
        if dynamic and zero:
            continue
        groups.append((p, samples, adv))
        if not dynamic and len(groups) == BATCH:
            break
    return groups, n_generated, n_zero


def ppo_update(model, groups, eps_low, eps_high, token_level):
    """
    用一批数据做 MU 次带裁剪的更新。

    【思路】先记下采样时（旧策略）每一步选那个做法的概率 π_old(a)；每次更新时用当前策略的概率算 ratio，
        按上面的裁剪规则决定这一步有没有梯度。
        归一化：sample-level（GRPO）= 每个回答内部除以它的步数，再除以回答的个数；
                token-level（DAPO）= 直接除以这一批所有回答的总步数。
    【返回】被裁剪掉的 (步数比例)，用来观察 clip 有多常起作用
    """
    data = []
    for p, samples, adv in groups:
        for s, a in zip(samples, adv):
            if a == 0.0:
                continue
            recs = step_records(p, s)
            old = [probs[action] for (key, action), probs in zip(recs, probs_for(model, p, s))]
            data.append((p, s, a, recs, old))
    if not data:
        return 0.0
    total_tokens = sum(len(recs) for _, _, _, recs, _ in data)
    n_responses = sum(len(samples) for _, samples, _ in groups)
    clipped, seen = 0, 0
    for _ in range(MU):
        grads = {}
        for p, s, a, recs, old in data:
            norm = 1.0 / total_tokens if token_level else 1.0 / (len(recs) * n_responses)
            for (key, action), probs, p_old in zip(recs, probs_for(model, p, s), old):
                ratio = probs[action] / p_old
                seen += 1
                if (a > 0 and ratio > 1 + eps_high) or (a < 0 and ratio < 1 - eps_low):
                    clipped += 1
                    continue
                add_grad(grads, key, action, probs, a * ratio * norm)
        apply_grads(model, grads, LR)
    return clipped / max(1, seen)


def train(train_set, test_set, seed, eps_high=0.2, dynamic=False, token_level=False):
    """
    训练 STEPS 个批次，返回学习曲线：
        list of (批次, 测试集 pass@1, 熵, 累计生成的样本数, 这一批里全 0 组的比例, 最近 10 批被裁剪的步数比例)。
    """
    rng = random.Random(seed)
    model = TrainableLLM(1e9)
    curve, generated, clip_fracs = [], 0, []
    for step in range(1, STEPS + 1):
        groups, n_gen, n_zero = collect_batch(model, train_set, rng, dynamic)
        generated += n_gen
        clip_fracs.append(ppo_update(model, groups, 0.2, eps_high, token_level))
        if step % 10 == 0:
            curve.append((step, evaluate(model, test_set, rng)[1], entropy_of(model, test_set[:60], rng),
                          generated, n_zero / max(1, n_gen // GROUP), mean(clip_fracs[-10:])))
    return model, curve


def main():
    train_set = make_problem_set(random.Random(160), n_problems=300, min_steps=2, max_steps=12)
    test_set = make_problem_set(random.Random(161), n_problems=200, min_steps=2, max_steps=12)
    configs = (("GRPO（ε = 0.2 / 0.2）", dict()),
               ("+ clip-higher（0.2 / 0.28）", dict(eps_high=0.28)),
               ("+ dynamic sampling", dict(eps_high=0.28, dynamic=True)),
               ("+ token-level loss（= DAPO）", dict(eps_high=0.28, dynamic=True, token_level=True)))
    curves = {name: train(train_set, test_set, seed=162, **kw)[1] for name, kw in configs}

    print("=" * 78)
    print(f"演示：GRPO 和 DAPO 的技巧（1B 模型，2~12 步题，每批 {BATCH} 组 × {GROUP} 个样本，每批更新 {MU} 次）")
    print("=" * 78)
    for title, idx, fmt in (("测试集 pass@1", 1, "{:.3f}"), ("策略的平均熵", 2, "{:.3f}"),
                            ("累计生成的样本数", 3, "{:d}"), ("这一批采到的组里，全对或全错的比例", 4, "{:.2f}"),
                            ("被裁剪（不贡献梯度）的步数比例", 5, "{:.3f}")):
        print(f"\n{title}：")
        print(pad("配置", 30) + "".join(pad(f"第 {s} 批", 10, "right") for s, *_ in curves[configs[0][0]]))
        for name, _ in configs:
            print(pad(name, 30) + "".join(pad(fmt.format(rec[idx]), 10, "right") for rec in curves[name]))

    print("\n观察（请对照上面的五张表）：")
    print("  · 全对或全错的组：训练越往后，模型越强，越来越多的题 8 个样本全对 —— 这些组的优势全为 0，")
    print("    对梯度没有任何贡献，却花掉了采样的成本。")
    print("  · dynamic sampling 丢掉这些组、凑满一批有信息的组：同样多的**更新次数**下学得快得多（第 30 批 0.958 vs 0.860），")
    print("    但多采了约 75% 的样本。按\"生成的样本数\"算，它的效率和 GRPO 差不多（第 20 批的 3864 个样本 → 0.859，")
    print("    GRPO 第 30 批的 3840 个样本 → 0.848）。它用采样的成本换来了每次更新的效率 —— 在真实的训练里，")
    print("    更新（反向传播）往往比采样更贵，所以这笔交易通常划算。")
    print("  · clip-higher 在这里几乎没有影响：裁剪只在不到 1% 的步骤上发生，放宽上界自然改变不了什么。")
    print("    真实的大规模训练里，低概率的 token 很容易在几次更新后就超出 1.2 的上界；DAPO 发现对称的裁剪会让熵")
    print("    很快崩塌（模型不再探索），放宽上界能保持多样性。这个 toy world 的参数太少、更新太温和，体现不出来。")
    print("  · token-level loss 在这里也没有带来明显的提升（这次略差，在噪声范围内；编写时换别的学习率试过，有时更好）。")
    print("    DAPO 的动机是长的推理链：sample-level 平均会让长回答里每个 token 的权重变小，")
    print("    长回答里的好模式学得慢、坏模式（重复、胡言乱语）也惩罚得不够。")
    print("  · 熵一路下降：策略越来越确定。toy world 里每一步只有一种正确做法，熵降到 0 也没关系。")
    print("  · 这些技巧在 toy world 里的效果大小，不代表在真实的大规模训练里的效果。DAPO 论文报告：")
    print("    用 Qwen2.5-32B 基础模型，四个技巧合起来在 AIME 2024 上达到 50 分，")
    print("    只用了 DeepSeek-R1-Zero-Qwen-32B（47 分）一半的训练步数。\n")


if __name__ == "__main__":
    main()