"""
=============================================================================
第 9 讲 · 共享模块：一个"玩具语言模型"
=============================================================================

【为什么需要一个玩具？】
    真正的 RLHF 要训练几十亿参数的语言模型，我们没法在笔记本上跑。
    但 RLHF 的核心问题——偏好数据、奖励模型、KL 约束、奖励过度优化、DPO——
    都可以在一个极简的模型上完整地展示出来，而且所有量都能精确计算。

    02、03、04 三个文件共用本模块。运行它们时当前目录必须是 code/lecture09/。

-----------------------------------------------------------------------------
【玩具设定】
-----------------------------------------------------------------------------
    - 只有一个"提示词"（prompt），模型从 N_RESP = 400 个候选"回答"中选一个。
      于是"语言模型"就是这 400 个回答上的一个概率分布 π(y)。
      （真实的语言模型是在天文数字个可能的回答上的分布，但数学结构完全一样。）

    - 每个回答 y 有两个特征：
        QUALITY[y] ：内在质量（标准正态分布）
        LENGTH[y]  ：长度（0~10，比如"段落数"）

    - **真实奖励**（代表人类真正的偏好，所有学习算法都看不到它的公式）：
        r*(y) = 质量 + 0.8·长度 - 0.08·长度²
      长度的贡献先增后减：长度 5 左右最好，太短（没说清楚）、太长（啰嗦）都不好。

    - **参考策略 π_ref**（相当于做完监督微调 SFT 的模型）：
        主要生成长度 1~4 的回答（以 2.5 为中心），并稍微偏好高质量的回答。
      在这个区间里，"更长"几乎总是"更好"——这为后面的"长度投机"埋下了伏笔。

    - **偏好数据**：从 π_ref 中抽两个回答，按 Bradley–Terry 模型（01 文件）和真实奖励
      决定人类更喜欢哪一个。这就是 RLHF 中"让人类比较模型的两个回答"。
"""

import numpy as np

N_RESP = 400
_rng = np.random.default_rng(0)
LENGTH = _rng.uniform(0, 10, N_RESP)
QUALITY = _rng.normal(0, 1, N_RESP)
FEATURES = np.stack([QUALITY, LENGTH], axis=1)          # 形状 (400, 2)


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def softmax(z):
    """数值稳定的 softmax（第 6 讲 01 文件）。"""
    z = z - np.max(z)
    e = np.exp(z)
    return e / e.sum()


def true_reward():
    """真实奖励 r*(y)。学习算法看不到它，只用来生成偏好标签和评价结果。"""
    return QUALITY + 0.8 * LENGTH - 0.08 * LENGTH ** 2


def reference_policy():
    """参考策略 π_ref（SFT 模型）：以长度 2.5 为中心的高斯形偏好，外加一点对质量的偏好。"""
    return softmax(-((LENGTH - 2.5) ** 2) / (2 * 1.2 ** 2) + 0.3 * QUALITY)


def kl(p, q):
    """KL(p ‖ q) = Σ p log(p / q)。p 中为 0 的项贡献为 0。"""
    mask = p > 0
    return float(np.sum(p[mask] * np.log(p[mask] / q[mask])))


def sample_preferences(policy, n_pairs, rng):
    """
    从 policy 中独立抽两个不同的回答，按 Bradley–Terry 模型和真实奖励生成偏好。

    【返回】(chosen, rejected)：两个整数数组，chosen[k] 是第 k 对中被人类偏好的回答编号。
        这就是 RLHF / DPO 文献中常见的 (y_w, y_l) 格式（w = win，l = lose）。
    """
    r = true_reward()
    i = rng.choice(N_RESP, size=n_pairs, p=policy)
    j = rng.choice(N_RESP, size=n_pairs, p=policy)
    keep = i != j
    i, j = i[keep], j[keep]
    i_wins = rng.random(len(i)) < sigmoid(r[i] - r[j])
    chosen = np.where(i_wins, i, j)
    rejected = np.where(i_wins, j, i)
    return chosen, rejected


def fit_linear_reward_model(chosen, rejected, lr=0.1, n_iters=3000):
    """
    训练线性奖励模型 r̂(y) = w · [质量, 长度]：Bradley–Terry 的极大似然（逻辑回归）。

        损失 = -mean( log σ(r̂(y_w) - r̂(y_l)) )

    【为什么是线性的？】为了简单、也为了让"外推出错"的机制一目了然：
        数据里的回答几乎都在长度 1~4 之间，在这个区间里真实奖励随长度上升，
        所以线性模型学到的是"越长越好"——它会把这个结论一直外推到长度 10。
        （神经网络奖励模型在远离训练数据的地方同样不可靠，只是出错的方式更难预测。）

    【返回】权重 w（长度 2）。
    """
    D = FEATURES[chosen] - FEATURES[rejected]
    w = np.zeros(FEATURES.shape[1])
    for _ in range(n_iters):
        p = sigmoid(D @ w)
        w += lr * D.T @ (1.0 - p) / len(p)          # 标签恒为 1（chosen 总是赢的那个）
    return w


def optimal_kl_policy(reward, ref, beta):
    """
    KL 正则 RL 目标的**闭式解**（讲义 3.3 节）：
        max_π  E_{y~π}[r(y)] - β KL(π ‖ π_ref)   的解是   π*(y) ∝ π_ref(y) · exp(r(y) / β)
    """
    return softmax(np.log(ref) + reward / beta)