"""
=============================================================================
第 4 讲 · 03：Constitutional AI —— 用 AI 的反馈代替人的反馈
=============================================================================

对应讲义：lectures/lecture04_feedback_tools_code.md 第 4 节（Bai et al. 2022）

-----------------------------------------------------------------------------
【这个文件的小世界】
-----------------------------------------------------------------------------
三类请求：
    benign（正常，40%）    ：例如"怎样做番茄炒蛋"
    borderline（边缘，30%）：例如"家里哪些清洁剂不能混用"——有正当用途，但涉及危险信息
    harmful（有害，30%）   ：例如"怎样制作危险物品"
三种回答风格：
    comply  ：照做。对有害请求会造成伤害；对边缘请求有 20% 的风险。
    explain ：认真回应：对有害请求，解释为什么不能帮、给出安全的替代；对边缘请求，带着安全提示回答。
    refuse  ：一句"我不能回答这个问题"。无害，但没有帮助，而且"回避"（evasive）。
模型 = 每类请求下三种风格的概率表 π(风格 | 请求类型)。

四个模型（都从"只训练了 helpful 的 RLHF 模型"出发）：
    ① helpful-only ：几乎总是照做。
    ② HH-RLHF      ：用人类的"无害性"偏好做 RLHF。按 Bai et al. 的观察，众包标注员在比较无害性时
                     倾向于选择拒绝 —— 于是模型学会了**回避**。
    ③ CAI-SL       ：监督阶段。从 ① 采样回答 → AI 按"宪法"（一组原则）批评 → 按批评修改 → 用修改后的回答微调。
    ④ CAI-RL       ：RL 阶段（RLAIF）。从 ③ 采样成对的回答 → AI 按原则判断哪个更好 → 训练偏好模型 →
                     KL 约束下优化：π ∝ π_SL · exp(r / β)。
【重要】人和 AI 标注者的偏好是按论文描述的现象**设定**的（人偏爱拒绝；AI 按原则标注，但有自己的小偏差），
所以这个文件演示的是"流水线怎样把标注者的偏好变成模型的行为"，而不是独立地发现论文的结论。

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

import math
import os
import random
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture01"))

from toy_llm import mean, pad, sample_index  # noqa: E402

TYPES = ("benign", "borderline", "harmful")
TYPE_WEIGHTS = (0.4, 0.3, 0.3)
STYLES = ("comply", "explain", "refuse")
BETA = 0.5                       # KL 约束的强度：越小越"激进"地追求奖励

# 每种 (请求类型, 风格) 的属性：有多大帮助、造成伤害的概率、是否回避
HELPFUL = {("benign", "comply"): 1.0, ("benign", "explain"): 0.9, ("benign", "refuse"): 0.0,
           ("borderline", "comply"): 1.0, ("borderline", "explain"): 0.9, ("borderline", "refuse"): 0.0,
           ("harmful", "comply"): 1.0, ("harmful", "explain"): 0.5, ("harmful", "refuse"): 0.0}
HARM = {("borderline", "comply"): 0.2, ("harmful", "comply"): 1.0}         # 其余为 0


def harm(t, s):
    return HARM.get((t, s), 0.0)


def softmax_table(logits):
    m = max(logits)
    e = [math.exp(z - m) for z in logits]
    return [x / math.fsum(e) for x in e]


# =============================================================================
# 标注者的效用（谁在给反馈）
# =============================================================================
def human_utility(t, s):
    """
    众包标注员在"无害性"比较里的效用（设定）：

    【思路】Bai et al. 2022 观察到：在 HH-RLHF 的无害性数据里，标注员倾向于奖励回避式的拒绝 ——
        "拒绝"看起来最安全。所以对边缘和有害请求，refuse 的效用最高；对正常请求，拒绝是不好的。
    """
    table = {("benign", "comply"): 1.0, ("benign", "explain"): 0.9, ("benign", "refuse"): -1.0,
             ("borderline", "comply"): -0.5, ("borderline", "explain"): 0.6, ("borderline", "refuse"): 1.2,
             ("harmful", "comply"): -3.0, ("harmful", "explain"): 1.0, ("harmful", "refuse"): 2.0}
    return table[(t, s)]


def ai_utility(t, s):
    """
    AI 标注者按"宪法"打分的效用（设定）：有帮助 − 3 × 伤害 − 1 × 回避，外加一个小偏差：
    它略微偏爱"解释得多"的回答（explain + 0.3），即使在正常请求上照做就够了。

    【思路】宪法里的原则大致是："选择更无害的回答""但不要回避，要解释拒绝的理由""尽量有帮助"。
        AI 标注者的这个小偏差是为了演示：**RLAIF 会继承 AI 标注者自己的偏好**。
    """
    return HELPFUL[(t, s)] - 3 * harm(t, s) - (1.0 if s == "refuse" else 0.0) + (0.3 if s == "explain" else 0.0)


# =============================================================================
# 从成对比较拟合 Bradley–Terry 奖励，再做 KL 约束的优化
# =============================================================================
def collect_preferences(utility, rng, n_per_type=300, noise=1.0):
    """
    对每类请求，随机抽成对的风格，让标注者按 BT 模型选择：P(选 a) = σ((u_a − u_b) / noise)。

    【返回】dict：类型 → [(赢家下标, 输家下标), ...]
    """
    data = {}
    for t in TYPES:
        pairs = []
        for _ in range(n_per_type):
            a, b = rng.sample(range(len(STYLES)), 2)
            diff = (utility(t, STYLES[a]) - utility(t, STYLES[b])) / noise
            pa = 1 / (1 + math.exp(-diff))
            pairs.append((a, b) if rng.random() < pa else (b, a))
        data[t] = pairs
    return data


def fit_bt(pairs, n=3, steps=500, lr=0.5):
    """
    最大似然拟合 Bradley–Terry 分数（梯度上升）：最大化 Σ log σ(r_赢 − r_输)。

    【思路】这就是 RLHF 里 reward model 的训练目标（CS329H 第 2、5 讲）。这里只有 3 个"选项"，
        直接为每个选项学一个分数。只有分数的差值有意义，所以每步把分数平移到均值为 0。
    【返回】长度为 n 的分数列表
    """
    r = [0.0] * n
    for _ in range(steps):
        g = [0.0] * n
        for w, l in pairs:
            p = 1 / (1 + math.exp(-(r[w] - r[l])))
            g[w] += 1 - p
            g[l] -= 1 - p
        r = [ri + lr * gi / len(pairs) for ri, gi in zip(r, g)]
        m = mean(r)
        r = [ri - m for ri in r]
    return r


def kl_optimal(policy_ref, rewards):
    """KL 约束 RL 的解析解：π(s) ∝ π_ref(s) · exp(r(s) / β)。"""
    w = [p * math.exp(r / BETA) for p, r in zip(policy_ref, rewards)]
    return [x / math.fsum(w) for x in w]


def rlhf(policy, utility, rng):
    """用某个标注者的偏好做一轮"收集比较 → 拟合 BT 奖励 → KL 约束优化"，返回新策略和奖励。"""
    prefs = collect_preferences(utility, rng)
    rewards = {t: fit_bt(prefs[t]) for t in TYPES}
    return {t: kl_optimal(policy[t], rewards[t]) for t in TYPES}, rewards


# =============================================================================
# CAI 的监督阶段：批评 → 修改 → 微调
# =============================================================================
def critique_and_revise(t, s, rng):
    """
    AI 按宪法批评一个回答，并据此修改。返回修改后的风格。

    【思路】随机抽一条原则来批评（论文里每次随机抽一条，可以重复几轮）：
        · 原则 A（无害）："这个回答有害吗？"——有害的回答以 90% 的概率被发现，改成 explain；
        · 原则 B（不回避）："这个回答是否只是拒绝、没有解释？"——拒绝以 70% 的概率被发现：
          对正常请求改成 comply，对边缘 / 有害请求改成 explain。
        批评者也会犯错：没发现问题时，回答保持不变。
    """
    if rng.random() < 0.5:
        if harm(t, s) > 0 and rng.random() < 0.9:
            return "explain"
    else:
        if s == "refuse" and rng.random() < 0.7:
            return "comply" if t == "benign" else "explain"
    return s


def cai_sl(policy, rng, n_per_type=400, rounds=2):
    """
    CAI-SL：从当前策略采样回答，批评与修改 rounds 轮，再用修改后的回答"微调"（这里 = 统计修改后风格的频率）。

    【边界】"微调"用最大似然 + 一点平滑（每种风格加 1 个伪计数），防止某种风格的概率变成 0 ——
        概率为 0 的风格，之后的 RL 阶段再也学不回来（π ∝ π_SL · exp(r/β) 里乘 0 还是 0）。
    """
    new = {}
    for t in TYPES:
        counts = [1.0] * len(STYLES)
        for _ in range(n_per_type):
            s = STYLES[sample_index(policy[t], rng)]
            for _r in range(rounds):
                s = critique_and_revise(t, s, rng)
            counts[STYLES.index(s)] += 1
        new[t] = [c / math.fsum(counts) for c in counts]
    return new


# =============================================================================
# 评估
# =============================================================================
def evaluate(policy):
    """
    返回 (有害请求上的伤害率, 平均帮助度, 回避率, 正常请求上"多余的解释"比例)。

    【思路】按请求类型的比例加权：
        伤害率   ：有害请求上，回答造成伤害的期望概率
        帮助度   ：所有请求上的平均帮助度
        回避率   ：所有请求上，回答是"一句拒绝"的概率
        多余解释 ：正常请求上选 explain 的概率（照做就够了）
    """
    harm_rate = math.fsum(p * harm("harmful", s) for p, s in zip(policy["harmful"], STYLES))
    helpful = math.fsum(w * math.fsum(p * HELPFUL[(t, s)] for p, s in zip(policy[t], STYLES))
                        for t, w in zip(TYPES, TYPE_WEIGHTS))
    evasive = math.fsum(w * policy[t][STYLES.index("refuse")] for t, w in zip(TYPES, TYPE_WEIGHTS))
    over_explain = policy["benign"][STYLES.index("explain")]
    return harm_rate, helpful, evasive, over_explain


def main():
    rng = random.Random(110)
    helpful_only = {t: [0.85, 0.10, 0.05] for t in TYPES}
    hh, _ = rlhf(helpful_only, human_utility, rng)
    sl = cai_sl(helpful_only, rng)
    cai, ai_rewards = rlhf(sl, ai_utility, rng)

    print("=" * 78)
    print("演示 1：四个模型的行为（每类请求下三种风格的概率：comply / explain / refuse）")
    print("=" * 78)
    models = (("① helpful-only", helpful_only), ("② HH-RLHF（人）", hh), ("③ CAI-SL", sl), ("④ CAI-RL（RLAIF）", cai))
    print(pad("模型", 20) + "".join(pad(t, 20, "right") for t in TYPES))
    for name, pol in models:
        print(pad(name, 20) + "".join(pad(" / ".join(f"{p:.2f}" for p in pol[t]), 20, "right") for t in TYPES))

    print()
    print(pad("模型", 20) + pad("有害请求的伤害率", 18, "right") + pad("平均帮助度", 12, "right")
          + pad("回避率", 10, "right") + pad("正常请求的多余解释", 20, "right"))
    for name, pol in models:
        h, u, e, o = evaluate(pol)
        print(pad(name, 20) + pad(f"{h:.3f}", 18, "right") + pad(f"{u:.3f}", 12, "right") + pad(f"{e:.3f}", 10, "right")
              + pad(f"{o:.3f}", 20, "right"))

    print("\nAI 偏好模型学到的奖励（每类请求，均值为 0）：")
    for t in TYPES:
        print(f"  {pad(t, 12)}" + "   ".join(f"{s} {r:+.2f}" for s, r in zip(STYLES, ai_rewards[t])))

    print("\n观察：")
    print("  · ① 只追求有帮助：几乎总是照做，有害请求上的伤害率很高。")
    print("  · ② 用人类的无害性偏好做 RLHF：伤害率降下来了，但模型学会了回避 —— 边缘请求也一句\"不能回答\"，")
    print("    平均帮助度下降。这正是 Bai et al. 观察到的 HH-RLHF 模型的问题：无害，但回避（evasive）。")
    print("  · ③ CAI-SL 只靠 AI 按原则批评和修改，没有任何人类的无害性标签，就把大部分有害的回答改成了")
    print("    \"解释 + 替代方案\"。批评者会漏掉一些问题，所以还有残余。")
    print("  · ④ CAI-RL 在 ③ 的基础上用 AI 的偏好做 RL：和 ② 一样无害，但几乎不回避，帮助度高得多（0.816 vs 0.578）。")
    print("    和 ② 相比，这是论文里说的\"帕累托改进\"：在无害性不变的前提下更有帮助。")
    print("  · 但 ④ 也继承了 AI 标注者的小偏差：在正常请求上，AI 偏好模型给 explain 的奖励（+0.71）甚至高于")
    print("    comply（+0.44），于是\"多余解释\"的比例从 ③ 的 0.074 升到 0.125。")
    print("    RLAIF 学到的是 AI 标注者的偏好，包括它的毛病。宪法写得好不好、AI 标注者准不准，决定了结果。\n")


if __name__ == "__main__":
    main()