"""
rlhf_lab.py —— RLHF 第 3 步与 DPO 的工具箱：序列概率、采样、KL 约束下的最优策略、REINFORCE / PPO / DPO（第 12 讲）
=====================================================================

【这个文件要回答的问题】
    第 11 讲做完了 RLHF 的前两步：指令微调得到 π_SFT（农场聊天助手），从人类排序里学了奖励模型 r_θ。
    第 3 步：让策略 π（一开始就是 π_SFT）去最大化
            E_{y ~ π(·|x)} [ r(x, y) ] − β · KL( π(·|x) ‖ π_ref(·|x) )             （π_ref = π_SFT）
    这个文件提供做这件事需要的所有零件，任务都是 describe（"describe : the dogs" → 一句关于狗的话）：
        - 一个回答的对数概率 log π(y | x) = Σ_t log π(y_t | x, y_<t)（含结尾的 </s>）；
        - 从策略采样；
        - 在一个有限的回答集合上精确算出 KL 约束下的最优策略 π*(y) ∝ π_ref(y) · exp(r(y)/β)；
        - REINFORCE、PPO、DPO 的训练循环。

【本文件提供】
    PROMPT_SUBJECTS / describe_prompt / random_prompts     提示
    seq_logprob / sample_responses                         序列概率与采样
    support_of / optimal_policy / policy_stats             有限支撑上的精确分析
    rollout_stats / train_pg                               REINFORCE 与 PPO
    dpo_loss / train_dpo                                   DPO
"""

import copy
import math
import os
import random
import sys

import torch
import torch.nn.functional as F

_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture11",):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from chat_lab import (ANIMALS, ASSISTANT, TEMPLATES, USER, generate, human_reward, load_chat_sft,  # noqa: E402
                      load_reward_model, noun_form, score)

PROMPT_SUBJECTS = [noun_form(a, pl) for a in ANIMALS for pl in (False, True)]       # 16 种主语：dog、dogs……


def describe_prompt(subj, template=("describe",)):
    """describe 任务的提示（词列表，不含 <s>）。"""
    return [USER] + list(template) + [":", "the", subj, ASSISTANT]


def random_prompts(rng, n):
    """n 个随机的 describe 提示（随机主语、随机的训练说法）。【返回】[(提示, 主语)]"""
    out = []
    for _ in range(n):
        subj = rng.choice(PROMPT_SUBJECTS)
        out.append((describe_prompt(subj, rng.choice(TEMPLATES["describe"][:4])), subj))
    return out


# =====================================================================
# 1. 序列概率与采样
# =====================================================================

def seq_logprob(model, vocab, prompts, responses):
    """
    log π(y | x)：回答里每个词（以及结尾的 </s>）的条件对数概率之和。
    【参数】prompts / responses：词列表的列表（responses 里如果有 "…" 表示没停下，会被去掉，也不加 </s>）
    【返回】形状 (N,) 的张量，带梯度（调用者需要时自己包 torch.no_grad()）
    【思路】一次前向算出整段的 logits（teacher forcing），只把回答部分的对数概率加起来。
        同样长度的放在一起算（GPT 用绝对位置，不能左边补齐），右边补 <pad>、用掩码去掉。
    """
    seqs, starts = [], []
    for p, r in zip(prompts, responses):
        stop = r[-1:] != ["…"]
        body = [w for w in r if w != "…"]
        ids = [vocab.bos_id] + vocab.encode(p) + vocab.encode(body) + ([vocab.eos_id] if stop else [])
        seqs.append(ids)
        starts.append(1 + len(p))
    T = max(len(s) for s in seqs)
    x = torch.full((len(seqs), T), vocab.pad_id, dtype=torch.long)
    mask = torch.zeros(len(seqs), T - 1)
    for i, (s, st) in enumerate(zip(seqs, starts)):
        x[i, :len(s)] = torch.tensor(s)
        mask[i, st - 1:len(s) - 1] = 1.0                 # 位置 t 的输出预测第 t+1 个词；回答从第 st 个词开始
    logits, _ = model(x[:, :-1])
    logp = torch.log_softmax(logits, dim=-1).gather(-1, x[:, 1:].unsqueeze(-1)).squeeze(-1)
    return (logp * mask).sum(1)


def sample_responses(model, vocab, prompts, gen, temperature=1.0):
    """从策略采样回答（温度 1 就是策略本身的分布）。"""
    return generate(model, vocab, prompts, greedy=False, gen=gen, temperature=temperature)


# =====================================================================
# 2. 有限支撑上的精确分析
# =====================================================================

@torch.no_grad()
def support_of(model, vocab, subj, n_samples, seed, batch=2000):
    """
    对一个提示，从策略里采 n_samples 个回答，把出现过的不同回答收集起来，当作"回答空间"的近似。
    【返回】(回答列表, 每个回答在模型下的精确 log π_ref(y|x) 张量, 这些回答的总概率)
    【边界】总概率 < 1：没采到的回答被忽略。采得越多越接近 1。
    """
    prompt = describe_prompt(subj)
    gen = torch.Generator().manual_seed(seed)
    seen = {}
    for s in range(0, n_samples, batch):
        for r in sample_responses(model, vocab, [prompt] * min(batch, n_samples - s), gen):
            if r[-1:] != ["…"]:
                seen[tuple(r)] = True
    resp = [list(r) for r in seen]
    logp = torch.cat([seq_logprob(model, vocab, [prompt] * len(resp[i:i + 500]), resp[i:i + 500])
                      for i in range(0, len(resp), 500)])
    return resp, logp, logp.exp().sum().item()


def optimal_policy(logp_ref, rewards, beta):
    """
    KL 约束下的最优策略（在给定的有限回答集合上）：π*(y) = π_ref(y) · exp(r(y)/β) / Z。
    【推导】max_π Σ_y π(y) r(y) − β Σ_y π(y) log(π(y)/π_ref(y))，约束 Σ π = 1。对 π(y) 求导令其为 0：
        r(y) − β (log π(y) − log π_ref(y) + 1) + λ = 0  ⇒  π(y) ∝ π_ref(y) exp(r(y)/β)。
    【参数】logp_ref：log π_ref（只需相对值，内部会在这个集合上归一化）；rewards：同长度的张量；beta > 0
    【返回】π*（在集合上归一化的概率张量）
    【边界】β → 0 时 π* 集中到奖励最高的那些回答上；β → ∞ 时 π* → π_ref。
    """
    z = logp_ref + rewards / beta
    return torch.softmax(z, dim=0)


def policy_stats(pi, logp_ref, true_r, rm_r, valid):
    """一个（集合上的）策略的期望真实奖励、期望奖励模型分数、KL(π ‖ π_ref)、说对话的比例。"""
    ref = torch.softmax(logp_ref, dim=0)
    kl = (pi * (torch.log(pi + 1e-30) - torch.log(ref + 1e-30))).sum().item()
    return {"true": (pi * true_r).sum().item(), "rm": (pi * rm_r).sum().item(), "kl": kl,
            "valid": (pi * valid).sum().item()}


# =====================================================================
# 3. REINFORCE 与 PPO
# =====================================================================

@torch.no_grad()
def rollout_stats(policy, ref, rm, vocab, n, seed):
    """
    用策略采 n 个回答（随机提示），统计：奖励模型分数、真实奖励 r*、"说对了"（r* > −2）的比例、
    KL(π ‖ π_ref) 的蒙特卡洛估计 E_π[log π(y|x) − log π_ref(y|x)]、没在 20 个词内停下的比例。
    """
    rng = random.Random(seed)
    gen = torch.Generator().manual_seed(seed)
    ps = random_prompts(rng, n)
    prompts = [p for p, _ in ps]
    resp = sample_responses(policy, vocab, prompts, gen)
    rm_s = score(rm, vocab, prompts, resp)
    true = [human_reward(s, r) for (_, s), r in zip(ps, resp)]
    lp = seq_logprob(policy, vocab, prompts, resp)
    lr = seq_logprob(ref, vocab, prompts, resp)
    return {"rm": sum(rm_s) / n, "true": sum(true) / n, "valid": sum(t > -2 for t in true) / n,
            "kl": (lp - lr).mean().item(), "nostop": sum(r[-1:] == ["…"] for r in resp) / n,
            "examples": [" ".join(r) for r in resp[:4]]}


def train_pg(policy, ref, vocab, reward_fn, steps, beta, algo="reinforce", batch=64, lr=1e-4, ppo_epochs=4,
             clip=0.2, seed=0, callback=None, every=0):
    """
    策略梯度：REINFORCE 或 PPO，奖励 = reward_fn(提示, 主语, 回答) − β · (log π(y|x) − log π_ref(y|x))。

    【参数】
        policy   : 要训练的策略（原地修改）；ref：参考策略 π_ref（冻结，用来算 KL 惩罚）
        reward_fn: f(prompts, subjects, responses) -> 奖励列表（奖励模型的分数，或者真实奖励 r*）
        beta     : KL 惩罚的系数（0 = 不加惩罚）
        algo     : "reinforce"——每批样本更新一次：loss = −mean((R − b) · log π(y|x))，b 是这批的平均奖励（基线）；
                   "ppo"——同一批样本更新 ppo_epochs 次，用重要性比率 ρ = π_新(y|x) / π_旧(y|x) 并裁剪到 [1−ε, 1+ε]：
                   loss = −mean( min(ρ·A, clip(ρ, 1−ε, 1+ε)·A) )，A = R − b
    【简化】真实的 PPO 在每个 token 上算优势（用一个价值网络和 GAE），我们把整句回答当作一个动作、用批平均做基线——
        slides 说"PPO 用轨迹相对于基线的优势"，这里的基线就是批平均。
    【返回】None；每 every 步调用一次 callback(step)
    """
    rng = random.Random(seed)
    gen = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    for step in range(1, steps + 1):
        ps = random_prompts(rng, batch)
        prompts, subjects = [p for p, _ in ps], [s for _, s in ps]
        policy.eval()
        with torch.no_grad():
            resp = sample_responses(policy, vocab, prompts, gen)
            old_lp = seq_logprob(policy, vocab, prompts, resp)
            ref_lp = seq_logprob(ref, vocab, prompts, resp)
        r = torch.tensor(reward_fn(prompts, subjects, resp), dtype=torch.float32)
        R = r - beta * (old_lp - ref_lp)
        A = R - R.mean()
        policy.train()
        for _ in range(ppo_epochs if algo == "ppo" else 1):
            lp = seq_logprob(policy, vocab, prompts, resp)
            if algo == "ppo":
                ratio = torch.exp(lp - old_lp)
                loss = -torch.min(ratio * A, torch.clamp(ratio, 1 - clip, 1 + clip) * A).mean()
            else:
                loss = -(A * lp).mean()
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            opt.step()
        policy.eval()
        if callback is not None and every and step % every == 0:
            callback(step)


# =====================================================================
# 4. DPO
# =====================================================================

def dpo_loss(policy, ref, vocab, prompts, winners, losers, beta):
    """
    DPO 的损失（Rafailov et al., 2023）：
        −log σ( β·[log π(y_w|x) − log π_ref(y_w|x)] − β·[log π(y_l|x) − log π_ref(y_l|x)] )
    【直觉】β·log(π/π_ref) 就是 DPO 的"隐式奖励"——把它代进 Bradley–Terry 模型，就是一个普通的二分类损失：
        提高赢的回答相对于参考模型的概率，降低输的。
    【返回】(损失, 赢的隐式奖励, 输的隐式奖励)
    """
    lw = seq_logprob(policy, vocab, prompts, winners)
    ll = seq_logprob(policy, vocab, prompts, losers)
    with torch.no_grad():
        rw = seq_logprob(ref, vocab, prompts, winners)
        rl = seq_logprob(ref, vocab, prompts, losers)
    zw, zl = beta * (lw - rw), beta * (ll - rl)
    return -F.logsigmoid(zw - zl).mean(), zw.detach(), zl.detach()


def train_dpo(policy, ref, vocab, pairs, beta, epochs=3, batch=64, lr=1e-4, seed=0, callback=None):
    """
    在偏好对 [(提示, 赢的回答, 输的回答)] 上训练 DPO（原地修改 policy）。callback(epoch) 每个 epoch 末调用。
    """
    rng = random.Random(seed)
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    order = list(range(len(pairs)))
    for ep in range(1, epochs + 1):
        rng.shuffle(order)
        policy.train()
        for s in range(0, len(order), batch):
            chunk = [pairs[i] for i in order[s:s + batch]]
            loss, _, _ = dpo_loss(policy, ref, vocab, [c[0] for c in chunk], [c[1] for c in chunk],
                                  [c[2] for c in chunk], beta)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            opt.step()
        policy.eval()
        if callback is not None:
            callback(ep)


NOISE = 0.5          # 标注员的噪声尺度，和第 11 讲 06 号脚本相同


def preference_data(sft, vocab, n_prompts, K, prompt_seed, sample_seed, label_seed):
    """
    造偏好数据：n_prompts 个 describe 提示，每个让 π_SFT（温度 1）采 K 个回答，由"标注员"排序
    （效用 = r* + NOISE·Gumbel 噪声），排序拆成 C(K,2) 个 (赢, 输) 对。
    【和第 11 讲的关系】做法和 lecture11/06_reward_model.py 完全一样；种子取 (1, 11, 21) 时得到的就是奖励模型的那份训练数据。
    【返回】(偏好对列表 [(提示, 赢的回答, 输的回答)], 提示列表 [(提示, 主语)])
    """
    rng = random.Random(prompt_seed)
    prompts = []
    for _ in range(n_prompts):
        a, pl = rng.choice(ANIMALS), rng.random() < 0.5
        tpl = rng.choice(TEMPLATES["describe"][:4])
        subj = noun_form(a, pl)
        prompts.append(([USER] + tpl + [":", "the", subj, ASSISTANT], subj))
    gen = torch.Generator().manual_seed(sample_seed)
    flat = [p for p, _ in prompts for _ in range(K)]
    outs = generate(sft, vocab, flat, greedy=False, gen=gen)
    responses = [outs[i * K:(i + 1) * K] for i in range(n_prompts)]
    lrng = random.Random(label_seed)
    pairs = []
    for (p, subj), rs in zip(prompts, responses):
        u = [human_reward(subj, r) - NOISE * math.log(-math.log(lrng.random() + 1e-12) + 1e-12) for r in rs]
        rank = sorted(range(len(rs)), key=lambda i: -u[i])
        for a_ in range(K):
            for b_ in range(a_ + 1, K):
                pairs.append((p, rs[rank[a_]], rs[rank[b_]], subj))
    return pairs, prompts


def load_models(log=None):
    """读取第 11 讲缓存的指令微调模型（π_SFT）和奖励模型。【返回】(sft, rm, vocab)"""
    sft, vocab = load_chat_sft(log=log)
    rm, _ = load_reward_model(log=log)
    for p in rm.parameters():
        p.requires_grad_(False)
    return sft, rm, vocab


def frozen_copy(model):
    m = copy.deepcopy(model)
    for p in m.parameters():
        p.requires_grad_(False)
    m.eval()
    return m


def entropy_bits(p):
    p = p[p > 0]
    return -(p * torch.log2(p)).sum().item()