"""
rlhf_world.py —— 第 18 讲：一个"人类偏好已知"的花园世界，用来研究奖励模型、RLHF 与 DPO
=====================================================================

【这个文件做什么】
    RLHF 的难点在于：我们想优化的东西（"有用、无害、让人满意"）没法写成损失函数，只能让人比较两个回答。
    真实研究里，人的偏好永远是黑箱；这里造一个**偏好已知**的版本：
        - "回答" = 一句花园句子（从 <bos> 生成到句号）；
        - "人"心里有一个**真实奖励** r*(y)（下面的 true_reward）：喜欢花、喜欢有形容词、喜欢"和朋友一起"，
          不太喜欢喂动物；**不合语法的句子**（语言模型被推得太远时会说出来）得 −2 分。
        - 人比较两个回答时，按 Bradley–Terry 模型给出偏好：P(y1 ≻ y2) = σ(r*(y1) − r*(y2))（有噪声，像真人一样）。
    因为花园世界只有约 3 万个合法句子，"在 KL 约束下最优的策略" π*(y) ∝ π_ref(y) exp(r(y)/β)
    可以**精确地**算出来——RLHF 和 DPO 想逼近的那个目标，在这里是看得见的。

【怎样使用】
    import rlhf_world as W
    base = W.base_policy()                          # "SFT 模型"：在花园句子上训练的语言模型
    ys = W.sample(base, 8, seed=0)                  # 采样 8 个回答
    W.true_reward(ys[0])                            # 真实奖励
    pairs = W.preference_pairs(base, 2000, seed=1)  # 2000 对带人类偏好标签的比较

【LLM 联系】InstructGPT（Ouyang et al. 2022）的三步：SFT → 奖励模型（人对回答排序）→ PPO + KL 惩罚；
            DPO（Rafailov et al. 2023）把后两步合成一步分类问题。
"""

import math
import os
import random
import sys

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import load_or_train, pad_batch, train  # noqa: E402
from toy_world import ADJECTIVES, NOUNS, GardenWorld, Vocab, garden_vocab  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")
WORLD = GardenWorld()
VOCAB = Vocab(garden_vocab().itos + ["<pad>"])
BOS, EOS, PAD = VOCAB.stoi["<bos>"], VOCAB.stoi["<eos>"], VOCAB.stoi["<pad>"]
DOT = VOCAB.stoi["."]
MAX_LEN = 14                                       # 回答最多 14 个 token（最长的合法句子是 11 个）

# =====================================================================
# 1. 真实奖励（"人心里的偏好"）
# =====================================================================

_CAT = {n: c for c, ns in NOUNS.items() for n in ns}
_ADJ = {a for adjs in ADJECTIVES.values() for a in adjs}


def true_reward(y):
    """
    真实奖励 r*(y)，y 是词列表（到句号为止）。

    【规则】不合语法：−2。合语法时从 0 开始：
            宾语是花 +1.0，是动物 −0.5；有形容词 +0.5；尾巴是 with … friends +0.5、from … own garden +0.3。
    【返回】float。合法句子的奖励在 −0.5 到 2.0 之间。
    【为什么这样设计】"不合语法 −2"模拟真实世界里"胡言乱语、重复、格式崩坏"的回答：人一眼就会讨厌它们，
                    但奖励模型只在正常的回答上训练过，没见过它们——这正是奖励被钻空子的地方。
    """
    y = list(y)
    if not WORLD.is_valid_sentence(y):
        return -2.0
    r = 0.0
    noun = next(w for w in y if w in _CAT)
    r += {"flower": 1.0, "animal": -0.5}.get(_CAT[noun], 0.0)
    r += 0.5 if any(w in _ADJ for w in y) else 0.0
    if "friends" in y:
        r += 0.5
    elif "garden" in y:
        r += 0.3
    return r


def all_sentences():
    """全部合法句子及其真实概率（第 1 讲的 enumerate_sentences）。【返回】(句子列表, 概率数组)。"""
    sents, probs = [], []
    for s, p in WORLD.enumerate_sentences():
        sents.append(list(s))
        probs.append(p)
    return sents, np.array(probs)


# =====================================================================
# 2. 策略（语言模型）：采样与 log 概率
# =====================================================================

ARCH = dict(d_model=64, n_heads=4, n_layers=2, d_ff=256)


def new_lm(seed=0):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="causal", **ARCH)


def base_policy(n=20000, epochs=4, seed=0):
    """
    "SFT 模型"（也是 RLHF 的参考策略 π_ref）：在 n 个按真实分布采样的花园句子上训练的语言模型。
    【边界】每个训练样例是一句话 + <eos>；生成时遇到句号就停。
    """
    rng = random.Random(seed)
    sents = [WORLD.sample_sentence(rng) for _ in range(n)]
    seqs = sorted(([BOS] + VOCAB.encode(s) for s in sents), key=len)
    batches = [pad_batch(seqs[i:i + 64], PAD) for i in range(0, len(seqs), 64)]

    def loss(m, b):
        ids, pad = b
        lg = m(ids[:, :-1], pad=pad[:, :-1])
        keep = ~pad[:, 1:]
        nll = F.cross_entropy(lg.reshape(-1, lg.shape[-1]), ids[:, 1:].reshape(-1), reduction="none")
        return (nll * keep.reshape(-1)).sum(), keep.sum()

    m, _ = load_or_train("base", {"n": n, "epochs": epochs, "seed": seed, "v": 1}, lambda: new_lm(seed),
                         lambda m: train(m, batches, loss, epochs=epochs, lr=3e-3, warmup=100, verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def sample(model, n, seed=0, temperature=1.0):
    """
    从策略采样 n 个回答（从 <bos> 开始，到句号为止，最多 MAX_LEN 个 token）。
    【返回】词列表的列表。没在 MAX_LEN 内说出句号的回答原样返回（必然不合语法）。
    """
    g = torch.Generator().manual_seed(seed)
    ids = torch.full((n, 1), BOS)
    done = torch.zeros(n, dtype=torch.bool)
    for _ in range(MAX_LEN):
        lg = model(ids)[:, -1] / temperature
        lg[:, [BOS, PAD, EOS]] = -1e9                 # 回答里不允许这三个 token
        nxt = torch.multinomial(torch.softmax(lg, -1), 1, generator=g)[:, 0]
        nxt = torch.where(done, torch.full_like(nxt, PAD), nxt)
        ids = torch.cat([ids, nxt[:, None]], 1)
        done |= nxt == DOT
        if done.all():
            break
    return [[VOCAB.itos[t] for t in row[1:] if t != PAD] for row in ids.tolist()]


def seq_logprob(model, ys, grad=False):
    """
    每个回答在策略下的 log 概率之和 log π(y) = Σ_t log π(y_t | y_<t)（和 sample 一样屏蔽 <bos>/<pad>/<eos>）。
    【参数】grad：True 时保留计算图（REINFORCE、DPO 要对它求梯度）。
    【返回】(n,) 张量。
    """
    seqs = [[BOS] + VOCAB.encode(y) for y in ys]
    ids, pad = pad_batch(seqs, PAD)
    ctx = torch.enable_grad() if grad else torch.no_grad()
    with ctx:
        lg = model(ids[:, :-1], pad=pad[:, :-1])
        mask = torch.zeros(lg.shape[-1], dtype=torch.bool)
        mask[[BOS, PAD, EOS]] = True
        lg = lg.masked_fill(mask, -1e9)
        lp = torch.log_softmax(lg, -1).gather(2, ids[:, 1:, None])[..., 0]
        return (lp * (~pad[:, 1:]).float()).sum(1)


# =====================================================================
# 3. 偏好数据与奖励模型
# =====================================================================


def preference_pairs(policy, n, seed=0):
    """
    从策略采样 n 对回答，"人"按 Bradley–Terry 给出偏好：P(y1 ≻ y2) = σ(r*(y1) − r*(y2))。
    【返回】[(更好的, 更差的)]。两个回答完全相同的对被丢掉重采（它们不提供任何信息）。
    【边界】偏好是有噪声的：r* 只差 0.5 时，"人"也有 38% 的时候选了更差的那个——真实的标注员之间的一致率
            通常只有 60%～75%（InstructGPT 报告约 73%）。
    """
    rng = random.Random(seed)
    out, k = [], 0
    while len(out) < n:
        ys = sample(policy, 2 * n, seed=seed * 1000 + k)
        k += 1
        for a, b in zip(ys[::2], ys[1::2]):
            if a == b:
                continue
            p = 1 / (1 + math.exp(-(true_reward(a) - true_reward(b))))
            out.append((a, b) if rng.random() < p else (b, a))
            if len(out) == n:
                break
    return out


class RewardModel(nn.Module):
    """
    奖励模型：一个 Transformer 读完回答，最后一个位置的隐藏状态经过线性层得到一个标量 r_φ(y)。
    InstructGPT 的奖励模型从 SFT 模型初始化（这里也从 base 初始化），只是把"预测下一个词"的头换成标量头。
    """

    def __init__(self, init=None, seed=0):
        super().__init__()
        self.lm = new_lm(seed)
        if init is not None:
            self.lm.load_state_dict(init.state_dict())
        self.head = nn.Linear(ARCH["d_model"], 1)

    def forward(self, ys):
        seqs = [[BOS] + VOCAB.encode(y) for y in ys]
        ids, pad = pad_batch(seqs, PAD)
        h = self.lm.hidden_states(ids, pad=pad)
        last = (~pad).sum(1) - 1
        return self.head(h[torch.arange(len(ys)), last])[:, 0]


def train_reward_model(name, pairs, init, epochs=3, lr=3e-4, seed=0):
    """
    Bradley–Terry 损失：−log σ(r_φ(y_w) − r_φ(y_l))，即"让更好的回答分数更高"的逻辑回归。
    【返回】训练好的 RewardModel（eval 模式）。
    """
    batches = [pairs[i:i + 64] for i in range(0, len(pairs), 64)]

    def loss(m, b):
        rw, rl = m([w for w, _ in b]), m([l for _, l in b])
        return -F.logsigmoid(rw - rl).sum(), torch.tensor(len(b))

    def build():
        return RewardModel(init, seed)

    m, _ = load_or_train(name, {"n": len(pairs), "epochs": epochs, "lr": lr, "seed": seed, "v": 1}, build,
                         lambda m: train(m, batches, loss, epochs=epochs, lr=lr, warmup=20, seed=seed, verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def rm_scores(rm, ys, batch=512):
    return torch.cat([rm(ys[i:i + batch]) for i in range(0, len(ys), batch)]).numpy()


# =====================================================================
# 4. KL 约束下的最优策略与 best-of-n（02 号脚本）
# =====================================================================


def reference_on_sentences(ref):
    """
    参考策略在全部 30336 个合法句子上的概率（重新归一化到合法句子上——base 模型约有 4% 的概率说出不合法的句子）。
    【返回】(句子, π_ref 概率, 真实奖励)。
    """
    S, _ = all_sentences()
    lp = torch.cat([seq_logprob(ref, S[i:i + 2048]) for i in range(0, len(S), 2048)]).numpy()
    p = np.exp(lp - lp.max())
    return S, p / p.sum(), np.array([true_reward(s) for s in S])


def kl_optimal(ref_p, R, beta):
    """
    KL 约束下的最优策略：max_π E_π[r] − β KL(π ‖ π_ref) 的解是 π*(y) = π_ref(y) exp(r(y)/β) / Z。
    【返回】(π*, E_π*[r], KL(π* ‖ π_ref))。
    """
    lw = np.log(ref_p + 1e-300) + R / beta
    lw -= lw.max()
    pi = np.exp(lw)
    pi /= pi.sum()
    kl = float((pi * (np.log(pi + 1e-300) - np.log(ref_p + 1e-300))).sum())
    return pi, float((pi * R).sum()), kl


def bon_expected(ref_p, R, n):
    """
    best-of-n（采样 n 个、留真实奖励最高的那个）的期望奖励，精确计算：P(最大值 ≤ v) = F(v)^n。
    【返回】float。对应的 KL 上界是 log n − (n − 1)/n（Stiennon et al. 2020 / Gao et al. 2023 用的公式）。
    """
    vals = np.unique(R)
    F_ = np.array([ref_p[R <= v].sum() for v in vals])
    F_prev = np.concatenate([[0.0], F_[:-1]])
    return float((vals * (F_ ** n - F_prev ** n)).sum())


def bon_kl(n):
    return math.log(n) - (n - 1) / n


# =====================================================================
# 5. 策略优化：REINFORCE + KL 惩罚（03 号脚本）与 DPO（04 号脚本）
# =====================================================================


def rl_train(name, ref, reward_fn, beta, steps=400, batch=64, lr=3e-4, seed=0, log_every=50, reward_tag=""):
    """
    最简单的 RLHF 策略优化：REINFORCE，奖励里减去 KL 惩罚。

    【每一步】
        1. 从当前策略 π_θ 采样 batch 个回答 y；
        2. 总奖励 R(y) = r(y) − β [log π_θ(y) − log π_ref(y)]（单样本的 KL 估计，只当作数值、不求导）；
        3. 优势 A = R − 这一批 R 的平均（基线，降低方差；GRPO 也是用同一组样本的平均作基线）；
        4. 损失 −A · log π_θ(y)，梯度就是策略梯度 E[A ∇ log π_θ(y)]。
    【参数】reward_fn：词列表的列表 → numpy 奖励；reward_tag：缓存键里区分不同奖励函数的名字。
    【返回】(策略, 训练曲线 [(步, 平均奖励, 平均真实奖励, 平均 KL, 合法比例)])。
    【LLM 联系】InstructGPT 用 PPO：多了价值网络、GAE、裁剪的重要性比率，让一批样本能安全地更新好几次；
                Ahmadian et al. (2024) 发现对 LLM，REINFORCE / RLOO 这类更简单的方法就够了，甚至更好。
    """
    cfg = {"beta": beta, "steps": steps, "batch": batch, "lr": lr, "seed": seed, "reward": reward_tag, "v": 1}
    curve = []

    def fit(pol):
        opt = torch.optim.Adam(pol.parameters(), lr=lr)
        for st in range(steps + 1):
            ys = sample(pol, batch, seed=seed * 1_000_003 + st)
            lp = seq_logprob(pol, ys, grad=True)
            kl = lp.detach() - seq_logprob(ref, ys)
            r = torch.tensor(reward_fn(ys), dtype=torch.float32)
            tot = r - beta * kl
            if st % log_every == 0:
                gold = np.array([true_reward(y) for y in ys])
                curve.append((st, float(r.mean()), float(gold.mean()), float(kl.mean()), float(np.mean(gold > -2))))
            if st == steps:
                break
            loss = -((tot - tot.mean()) * lp).mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
        return curve

    def build():
        m = new_lm()
        m.load_state_dict(ref.state_dict())
        return m

    pol, info = load_or_train(f"rl_{name}", cfg, build, fit, verbose=False, ckpt_dir=CKPT_DIR)
    pol.eval()
    return pol, info["history"]


def dpo_train(name, ref, pairs, beta, epochs=3, lr=1e-4, batch=64, seed=0):
    """
    DPO（Rafailov et al. 2023）：
        L = −log σ( β [ (log π_θ(y_w) − log π_ref(y_w)) − (log π_θ(y_l) − log π_ref(y_l)) ] )
    就是把"隐式奖励" r_θ(y) = β log π_θ(y)/π_ref(y) 代入奖励模型的 Bradley–Terry 损失。
    不需要奖励模型，也不需要从策略采样——是在固定的偏好数据上的分类问题。
    【返回】训练好的策略。
    """

    def fit(pol):
        opt = torch.optim.Adam(pol.parameters(), lr=lr)
        rw = seq_logprob(ref, [w for w, _ in pairs])
        rl = seq_logprob(ref, [l for _, l in pairs])
        g = torch.Generator().manual_seed(seed)
        for _ in range(epochs):
            perm = torch.randperm(len(pairs), generator=g).tolist()
            for i in range(0, len(pairs), batch):
                idx = perm[i:i + batch]
                lw = seq_logprob(pol, [pairs[j][0] for j in idx], grad=True)
                ll = seq_logprob(pol, [pairs[j][1] for j in idx], grad=True)
                loss = -F.logsigmoid(beta * ((lw - rw[idx]) - (ll - rl[idx]))).mean()
                opt.zero_grad()
                loss.backward()
                opt.step()
        return {}

    def build():
        m = new_lm()
        m.load_state_dict(ref.state_dict())
        return m

    cfg = {"n": len(pairs), "beta": beta, "epochs": epochs, "lr": lr, "seed": seed, "v": 1}
    pol, _ = load_or_train(f"dpo_{name}", cfg, build, fit, verbose=False, ckpt_dir=CKPT_DIR)
    pol.eval()
    return pol


def evaluate_policy(pol, ref, n=1000, seed=123):
    """
    【返回】dict：真实奖励的平均、KL(π ‖ π_ref) 的估计、合法句子的比例、不同回答的比例（多样性）、样例。
    """
    ys = sample(pol, n, seed=seed)
    g = np.array([true_reward(y) for y in ys])
    kl = float((seq_logprob(pol, ys) - seq_logprob(ref, ys)).mean())
    return {"gold": float(g.mean()), "kl": kl, "gram": float(np.mean(g > -2)),
            "distinct": len(set(map(tuple, ys))) / n, "samples": ys[:4]}


__all__ = ["WORLD", "VOCAB", "true_reward", "all_sentences", "new_lm", "base_policy", "sample", "seq_logprob",
           "preference_pairs", "RewardModel", "train_reward_model", "rm_scores", "MAX_LEN", "reference_on_sentences",
           "kl_optimal", "bon_expected", "bon_kl", "rl_train", "dpo_train", "evaluate_policy"]