"""
reason_lab.py —— 第 20 讲的共享模块：一个"会选择想多久"的小推理模型，以及测试时计算、STaR、GRPO 要用的工具
=====================================================================

【这个文件要回答的问题】
    推理模型（o1、DeepSeek-R1）的两个旋钮：**训练时**用强化学习教它"想"，**测试时**让它想得更久、采样更多次。
    要在笔记本上观察这两件事，需要一个任务：
        1. 不想（直接答）很难，想得越完整越容易答对；
        2. 模型自己决定想多少（就像 R1 自己决定思考链的长度）；
        3. 答案对不对可以用规则检查（R1-Zero 的"基于规则的奖励"），中间步骤也可以逐步检查（过程验证）。

【任务：数字和取模 10】
    题目：k 个数字 d₁ … d_k（k = 2…8），答案 = (d₁ + … + d_k) mod 10。
    模型在 "=" 之后先写一个**格式记号**，然后：
        直接回答 "A"：   <s> 3 5 9 2 7 _ _ _ = A 6 </s>
        简略思考 "S"：   <s> 3 5 9 2 7 _ _ _ = S 8 9 6 A 6 </s>        每两个数字写一次累加和（mod 10）：要一步加两个数
        完整思考 "T"：   <s> 3 5 9 2 7 _ _ _ = T 3 8 7 9 6 A 6 </s>    每个数字写一次累加和：每一步只加一个数
    （题目写在固定宽度的 8 个格子里，空的格子是 "_"。）
    思考越完整，每一步越简单（第 11 讲 §4.2：每个 token 只有固定层数的计算），但要生成的 token 越多。
    "预训练"的最后阶段三种格式的比例是 50% / 30% / 20%：模型学会了三种写法，但更常"偷懒"。

【本文件提供】
    词表与格式：VOCAB、ID、FORMATS、make_digits、solution、encode、parse、check_chain
    数据与模型：batch_fn（按格式比例造一批训练数据）、build_model（第 2 讲的 GPT，4 层）、BASE_CFG（基础模型的两阶段配方）、
        get_model（训练并缓存）
    生成与评估：generate（批量采样 / 贪心，可以强制格式）、evaluate、seq_logprobs（每个序列每个 token 的对数概率）
    强化学习：rl_train（GRPO / 不带基线的 REINFORCE；组内标准化的优势、裁剪的比率、k3 估计的 KL 惩罚）

【边界】
    - "模型"只有约 20 万参数，任务只是加法；结论是"机制上会发生什么"，不是大模型上的定量结果。
    - 生成时每一步都把整个序列重新喂给模型（序列不超过 24 个 token，不需要 KV 缓存）。
"""

import hashlib
import os
import random
import sys

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
_p = os.path.join(HERE, "..", "lecture02")
if _p not in sys.path:
    sys.path.insert(0, _p)

from lm_models import GPT, train_on_batches  # noqa: E402

CKPT = os.path.join(HERE, "checkpoints")
VOCAB = ["<pad>", "<s>", "</s>", "=", "A", "S", "T", "_"] + [str(i) for i in range(10)]
ID = {w: i for i, w in enumerate(VOCAB)}
DIGIT0 = ID["0"]
FORMATS = {"direct": "A", "short": "S", "full": "T"}
K_MIN, K_MAX = 2, 8
MAX_LEN = 24
IGNORE = -100


# =====================================================================
# 1. 题目、三种写法、解析与检查
# =====================================================================

def make_digits(rng, k):
    return [rng.randrange(10) for _ in range(k)]


def solution(digits, fmt):
    """
    "=" 之后的正确输出（词的列表）。
    【例】digits = [3, 5, 9, 2, 7]：
        direct → ["A", "6", "</s>"]
        short  → ["S", "8", "9", "6", "A", "6", "</s>"]     累加和写在第 2、4、5 个数字之后（k 为奇数时最后一步只加一个）
        full   → ["T", "3", "8", "7", "9", "6", "A", "6", "</s>"]
    """
    ans = sum(digits) % 10
    if fmt == "direct":
        return ["A", str(ans), "</s>"]
    step = 1 if fmt == "full" else 2
    chain, s = [], 0
    for i, d in enumerate(digits, start=1):
        s = (s + d) % 10
        if i % step == 0 or i == len(digits):
            chain.append(str(s))
    return [FORMATS[fmt]] + chain + ["A", str(ans), "</s>"]


def prompt_ids(digits):
    """
    提示：<s> + 数字 + 空格 "_"（补到固定的 K_MAX 格）+ "="。
    【为什么补成固定宽度】像填表一样，题目总是写在 8 个格子里。这样写第 t 个中间结果时，要找的第 t 个数字
        永远在固定的距离之外。第一版不补齐，"往回多远"随题目长度 k 变化，小模型在几千步里常常学不会（README 第 2 条）。
    """
    return [ID["<s>"]] + [DIGIT0 + d for d in digits] + [ID["_"]] * (K_MAX - len(digits)) + [ID["="]]


def encode(digits, fmt):
    """一道题的完整序列（整数编号），以及从哪个位置开始计损失（"=" 之后的所有词，包括格式记号）。"""
    p = prompt_ids(digits)
    return p + [ID[w] for w in solution(digits, fmt)], len(p) - 1


def parse(out_ids):
    """
    解析 "=" 之后生成的编号列表。
    【返回】字典：fmt（"direct" / "short" / "full" / None）、chain（中间结果的数字列表）、answer（整数或 None）、
        ok（格式是否完整：格式记号 + 若干数字 + A + 一个数字 + </s>，且直接回答没有中间步骤、思考有中间步骤）、
        length（到 </s> 为止生成了几个词）
    【边界】没有 "A 数字" 时 answer = None；有答案但格式不完整（例如没写 </s>）时 answer 照常给出、ok = False。
    """
    words = [VOCAB[i] for i in out_ids]
    if "</s>" in words:
        words = words[:words.index("</s>") + 1]
    res = {"fmt": None, "chain": [], "answer": None, "ok": False, "length": len(words)}
    inv = {v: k for k, v in FORMATS.items()}
    if not words or words[0] not in inv:
        return res
    res["fmt"] = inv[words[0]]
    i = 1
    if res["fmt"] != "direct":                       # 直接回答的格式记号 "A" 本身就是答案标记，后面没有中间步骤
        while i < len(words) and words[i].isdigit():
            res["chain"].append(int(words[i]))
            i += 1
        if i < len(words) and words[i] == "A":
            i += 1
        else:
            return res
    if i < len(words) and words[i].isdigit():
        res["answer"] = int(words[i])
        res["ok"] = (i + 1 < len(words) and words[i + 1] == "</s>"
                     and (res["fmt"] == "direct") == (len(res["chain"]) == 0))
    return res


def check_chain(digits, parsed):
    """
    过程验证：中间结果每一步都对、答案等于最后一步的累加和吗？（不需要知道标准答案，只用题目本身就能检查）
    【返回】True / False；直接回答没有中间步骤，无法验证，返回 None
    """
    if parsed["fmt"] is None or not parsed["ok"]:
        return False
    if parsed["fmt"] == "direct":
        return None
    want = solution(digits, parsed["fmt"])
    chain = [int(w) for w in want[1:-3]]
    return parsed["chain"] == chain and parsed["answer"] == int(want[-2])


# =====================================================================
# 2. 训练数据与模型
# =====================================================================

def batch_fn(seed, batch, mix=(0.5, 0.3, 0.2), k_range=(K_MIN, K_MAX), fixed=None, fixed_frac=1.0):
    """
    造一批训练数据：(inputs, targets)，形状都是 (batch, MAX_LEN − 1)。
    【参数】mix：直接 / 简略 / 完整三种格式的比例；k_range：题目长度的范围；
        fixed：给一个 [(digits, "=" 之后整段输出的词列表), ...] 的列表时，从里面有放回地抽（STaR 的数据）；
        fixed_frac：每个样本以这个概率从 fixed 里抽，否则按 mix 现造一道新题（"少量人写的思考链 + 大量新题"）
    【损失】只在 "=" 之后的词上计算（格式记号也算：模型要学会"选择"格式的分布）。
    """
    rng = random.Random(seed)
    inp = torch.full((batch, MAX_LEN - 1), ID["<pad>"], dtype=torch.long)
    tgt = torch.full((batch, MAX_LEN - 1), IGNORE, dtype=torch.long)
    for b in range(batch):
        if fixed is not None and rng.random() < fixed_frac:
            digits, out = fixed[rng.randrange(len(fixed))]
            ids = prompt_ids(digits) + [ID[w] for w in out]
            start = K_MAX + 1
        else:
            digits = make_digits(rng, rng.randint(*k_range))
            fmt = rng.choices(list(FORMATS), mix)[0]
            ids, start = encode(digits, fmt)
        n = len(ids) - 1
        inp[b, :n] = torch.tensor(ids[:-1])
        tgt[b, start:n] = torch.tensor(ids[start + 1:])
    return inp, tgt


def build_model(seed=0):
    """第 2 讲的 GPT：d = 64、4 层、4 个头、可学习的位置嵌入（最长 24）。构造前先设种子（本机上每个进程的默认种子不同）。"""
    torch.manual_seed(seed)
    return GPT(len(VOCAB), d_model=64, n_layers=4, n_heads=4, max_len=MAX_LEN)


BASE_CFG = {"seed": 1, "lr": 1e-3,
            "phases": [[[0.0, 0.5, 0.5], 2000, 100, [K_MIN, K_MAX]], [[0.5, 0.3, 0.2], 1500, 50, [K_MIN, K_MAX]]]}
"""
基础模型的配方（"预训练"）：两个阶段，每个是 [格式比例, 步数, 预热步数, 题目长度范围]。
    阶段 1：只有简略思考和完整思考（各一半），2000 步——先学会"一步一步加"；
    阶段 2：直接 / 简略 / 完整 = 50% / 30% / 20%，1500 步——学会三种写法，并且更常"偷懒"。
【为什么分两个阶段、为什么种子是 1】三种格式从一开始就混在一起训练时，4000–5000 步都停在平台上（只会 k = 2）；
    先只练思考链就学会了。种子 0 在两阶段配方下也只学会了完整思考（简略思考没学会），种子 1 三种都学会了（README 第 3、4 条）。
"""


def get_model(name, cfg, quick=False, log=print, init=None):
    """
    有缓存就读取，没有就训练并缓存到 checkpoints/<name>.pt。
    【参数】cfg：{"seed", "lr", "phases": [[格式比例, 步数, 预热, [k_min, k_max]], ...], 可选 "fixed"（STaR 的数据）}；
        init：从这个模型的权重开始训练（STaR 每一轮都从基础模型重新开始）
    【边界】quick=True 时每个阶段只训练 30 步、不读写缓存；读取时配置不同就重新训练。
    """
    path = os.path.join(CKPT, f"{name}.pt")
    key = {k: v for k, v in cfg.items() if k != "fixed"}
    key["n_fixed"] = len(cfg.get("fixed") or [])
    # 微调数据（STaR 生成的思考链）的指纹：数据变了，缓存就作废
    key["fixed_hash"] = hashlib.sha256(repr(cfg.get("fixed")).encode()).hexdigest()[:16]
    if not quick and os.path.exists(path):
        blob = torch.load(path, weights_only=True)
        if blob["cfg"] == key:
            model = build_model(cfg.get("seed", 0))
            model.load_state_dict(blob["state"])
            model.eval()
            if log:
                log(f"  （读取缓存的模型：checkpoints/{name}.pt）")
            return model
    model = build_model(cfg.get("seed", 0))
    if init is not None:
        model.load_state_dict(init.state_dict())
    data_seed = cfg.get("seed", 0) * 1_000_003
    for i, (mix, steps, warmup, k_range) in enumerate(cfg["phases"]):
        steps = 30 if quick else steps
        train_on_batches(model, lambda s, i=i, mix=mix, k_range=k_range: batch_fn(
            data_seed + i * 10 ** 6 + s, 64, mix=tuple(mix), k_range=tuple(k_range), fixed=cfg.get("fixed"),
            fixed_frac=cfg.get("fixed_frac", 1.0)),
            steps, lr=cfg.get("lr", 1e-3), log_every=0, log=None, warmup=min(warmup, steps))
    model.eval()
    if not quick:
        os.makedirs(CKPT, exist_ok=True)
        torch.save({"cfg": key, "state": model.state_dict()}, path)
        if log:
            log(f"  （已保存到 checkpoints/{name}.pt）")
    return model


# =====================================================================
# 3. 生成与评估
# =====================================================================

@torch.no_grad()
def generate(model, prompts, gen=None, temperature=1.0, force=None, max_new=16):
    """
    批量生成 "=" 之后的输出。
    【参数】prompts：每个是提示的编号列表（以 "=" 结尾）；gen：torch.Generator（None 或 temperature=0 时贪心）；
        force：强制的开头（例如 ["T"] 表示"必须完整思考"——测试时控制想多久）
    【返回】输出编号的列表的列表（不含提示；到 </s> 为止，没写 </s> 就截在 max_new 个）
    【边界】右侧填充；可学习位置嵌入 + 因果掩码，填充不影响前面的位置。
    """
    model.eval()
    seqs = [list(p) + ([ID[w] for w in force] if force else []) for p in prompts]
    starts = [len(p) for p in prompts]
    done = [False] * len(seqs)
    for _ in range(max_new):
        alive = [i for i in range(len(seqs)) if not done[i] and len(seqs[i]) < MAX_LEN]
        if not alive:
            break
        T = max(len(seqs[i]) for i in alive)
        ids = torch.full((len(alive), T), ID["<pad>"], dtype=torch.long)
        for r, i in enumerate(alive):
            ids[r, :len(seqs[i])] = torch.tensor(seqs[i])
        logits, _ = model(ids)
        last = torch.stack([logits[r, len(seqs[i]) - 1] for r, i in enumerate(alive)])
        if gen is None or temperature == 0:
            nxt = last.argmax(-1)
        else:
            nxt = torch.multinomial(F.softmax(last / temperature, -1), 1, generator=gen).squeeze(1)
        for r, i in enumerate(alive):
            seqs[i].append(int(nxt[r]))
            if int(nxt[r]) == ID["</s>"]:
                done[i] = True
        for i in range(len(seqs)):
            if len(seqs[i]) >= MAX_LEN:
                done[i] = True
    return [s[st:] for s, st in zip(seqs, starts)]


def evaluate(model, problems, gen=None, temperature=0.0, force=None):
    """
    【返回】字典：acc（答案正确的比例）、length（平均生成长度）、share（三种格式的比例）、
        chain_ok（有思考链的回答里，中间步骤全对的比例）、results（每道题的 parse 结果）
    """
    outs = generate(model, [prompt_ids(d) for d in problems], gen, temperature, force)
    res = [parse(o) for o in outs]
    acc = sum(r["answer"] == sum(d) % 10 for r, d in zip(res, problems)) / len(problems)
    share = {f: sum(r["fmt"] == f for r in res) / len(res) for f in FORMATS}
    chains = [check_chain(d, r) for r, d in zip(res, problems) if r["fmt"] in ("short", "full")]
    return {"acc": acc, "length": sum(r["length"] for r in res) / len(res), "share": share,
            "chain_ok": sum(bool(c) for c in chains) / len(chains) if chains else float("nan"), "results": res,
            "outputs": outs}


def seq_logprobs(model, prompts, outputs):
    """
    每个序列的每个输出 token 的对数概率 log π(y_t | x, y_<t)。
    【返回】(logp (B, L), mask (B, L))：L = 最长的输出长度；mask = 1 表示这个位置是真实的输出 token
    """
    full = [p + o for p, o in zip(prompts, outputs)]
    T = max(len(s) for s in full)
    ids = torch.full((len(full), T), ID["<pad>"], dtype=torch.long)
    for b, s in enumerate(full):
        ids[b, :len(s)] = torch.tensor(s)
    logits, _ = model(ids[:, :-1])
    lp = F.log_softmax(logits, -1).gather(-1, ids[:, 1:].unsqueeze(-1)).squeeze(-1)
    L = max(len(o) for o in outputs)
    out = torch.zeros(len(full), L)
    mask = torch.zeros(len(full), L)
    for b, (p, o) in enumerate(zip(prompts, outputs)):
        out[b, :len(o)] = lp[b, len(p) - 1:len(p) - 1 + len(o)]
        mask[b, :len(o)] = 1
    return out, mask


# =====================================================================
# 4. 强化学习：GRPO
# =====================================================================

def rule_reward(digits, parsed, format_bonus=0.1):
    """
    R1-Zero 式的基于规则的奖励：准确率奖励（答案对得 1）+ 格式奖励（输出完整地遵守模板得 format_bonus）。
    没有奖励模型——答案对不对用规则算，所以没有"奖励被钻空子"的问题（第 12 讲 §3.4）。
    """
    r = 1.0 if parsed["answer"] == sum(digits) % 10 else 0.0
    return r + (format_bonus if parsed["ok"] else 0.0)


def rl_train(policy, ref, steps, algo="grpo", prompts_per_step=16, group=8, lr=3e-5, beta=0.04, clip=0.2,
             inner=2, seed=0, k_range=(K_MIN, K_MAX), callback=None, every=0):
    """
    GRPO（Shao et al., 2024，DeepSeekMath；R1 用的算法）或不带基线的 REINFORCE。

    【GRPO 的一步】
        1. 抽 prompts_per_step 道题，每道题用当前策略采样 group 个回答（温度 1）；
        2. 用规则算每个回答的奖励 r_i；**组内标准化**得到优势 A_i = (r_i − mean(r)) / (std(r) + ε)——
           同一道题的几个回答互相当基线，不需要 PPO 那样的价值网络（slides："去掉价值模型，大大减少显存"）；
        3. 回答里的每个 token 都用这个优势；目标是 PPO 式的裁剪比率：
               min(ρ_t A, clip(ρ_t, 1−ε, 1+ε) A)，ρ_t = π_θ(y_t|…) / π_old(y_t|…)
           再减去 β·KL(π_θ ‖ π_ref)，KL 用 DeepSeekMath 的无偏估计 k3 = π_ref/π_θ − log(π_ref/π_θ) − 1（逐 token）；
        4. 同一批样本更新 inner 次（第一次 ρ = 1，之后 ρ ≠ 1，裁剪才起作用）。
    【algo="reinforce"】对照：优势就是奖励本身 A_i = r_i（不减基线），不裁剪、不加 KL——第 12 讲的 REINFORCE 去掉基线。
    【参数】policy（原地修改）、ref（冻结的参考模型）、steps、lr、beta、clip、inner、k_range（题目长度）
    【返回】每步的记录列表：{"step", "reward", "acc", "length", "share", "kl"}；每 every 步调用 callback(step, policy)
    """
    rng = random.Random(seed)
    gen = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    hist = []
    for step in range(1, steps + 1):
        probs = [make_digits(rng, rng.randint(*k_range)) for _ in range(prompts_per_step)]
        prompts = [prompt_ids(d) for d in probs for _ in range(group)]
        digits = [d for d in probs for _ in range(group)]
        outs = generate(policy, prompts, gen, temperature=1.0)
        parsed = [parse(o) for o in outs]
        r = torch.tensor([rule_reward(d, p) for d, p in zip(digits, parsed)])
        if algo == "grpo":
            g = r.view(prompts_per_step, group)
            A = ((g - g.mean(1, keepdim=True)) / (g.std(1, keepdim=True) + 1e-4)).view(-1)
        else:
            A = r.clone()
        with torch.no_grad():
            old_lp, mask = seq_logprobs(policy, prompts, outs)
            ref_lp, _ = seq_logprobs(ref, prompts, outs)
        policy.train()
        for _ in range(inner if algo == "grpo" else 1):
            lp, _ = seq_logprobs(policy, prompts, outs)
            if algo == "grpo":
                ratio = torch.exp(lp - old_lp)
                surr = torch.min(ratio * A[:, None], torch.clamp(ratio, 1 - clip, 1 + clip) * A[:, None])
                kl = torch.exp(ref_lp - lp) - (ref_lp - lp) - 1
                per_tok = surr - beta * kl
            else:
                per_tok = lp * A[:, None]
            # 每个回答先在自己的 token 上取平均，再在回答之间取平均（GRPO 论文的写法）
            loss = -((per_tok * mask).sum(1) / mask.sum(1)).mean()
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            opt.step()
        policy.eval()
        with torch.no_grad():
            kl_now = ((torch.exp(ref_lp - old_lp) - (ref_lp - old_lp) - 1) * mask).sum() / mask.sum()
        acc = sum(p["answer"] == sum(d) % 10 for p, d in zip(parsed, digits)) / len(digits)
        share = {f: sum(p["fmt"] == f for p in parsed) / len(parsed) for f in FORMATS}
        hist.append({"step": step, "reward": r.mean().item(), "acc": acc,
                     "length": sum(p["length"] for p in parsed) / len(parsed), "share": share, "kl": kl_now.item()})
        if callback is not None and every and step % every == 0:
            callback(step, policy)
    return hist