"""
=============================================================================
第 2 讲 · 代码 2：五种偏好数据，一个框架 —— 以及"哪种问法最划算"
=============================================================================

对应讲义：lectures/lecture02_choice_data_bt.md 第 3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
收集偏好数据时，第一个要做的决定是"怎么问"：
    · 给两个回答，问哪个好？（两两比较：Chatbot Arena、大多数 RLHF 数据）
    · 给四个回答，挑最好的？（多选一：推荐页面上的点击）
    · 给四个回答，从好到坏排个序？（排序：InstructGPT）
    · 给一个回答，问能不能接受？（接受 / 拒绝：点赞 / 点踩、"这个回答可以发布吗"）
这个文件用第 1 讲的写作助手世界（专家标注员）做四个实验：

    演示 1：五种数据长什么样，怎样都拆成"从集合中选一个"的事件。
    演示 2：一个拟合器搞定所有形式 —— 数据足够多时，每种形式都能把效用学回来；
            把排序拆成两两比较再用 BT 拟合，结果和 Plackett–Luce 几乎一样（这不是巧合）。
    演示 3：同样的标注成本，哪种问法学得最好？—— 按"读了多少个回答"计成本，比较估计误差。
    演示 4：接受 / 拒绝数据能告诉我们"绝对水平"，两两比较不能 —— 两个世界的比较数据一模一样，
            一个世界里的回答全都很差。
    演示 5：把排序拆成的两两比较**不是独立的**：当成独立数据，会低估不确定性。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_data_formats.py        # 约 10 秒

只用到 Python 标准库，以及第 1 讲的 toy_world.py、本目录的 choice_data.py。
"""

import math
import random

from choice_data import (OUTSIDE, accept_reject_to_event, choice_to_event, fit_choice_model, pair_to_event,
                         ranking_to_events, ranking_to_pairs, sample_accept, sample_choice, sample_pair,
                         sample_ranking)
from toy_world import KEYS, kendall_tau, make_expert, mean, pad, print_header, BY_KEY

EXPERT = make_expert()
TRUE = {k: EXPERT.utility(BY_KEY[k]) for k in KEYS}             # 专家心里的效用（上帝视角）
TRUE_CENTERED = {k: TRUE[k] - mean(list(TRUE.values())) for k in KEYS}
THRESHOLD = 3.5        # 接受 / 拒绝时外部选项的效用："至少要这么好，我才愿意接受"


def centered(scores):
    """把分数平移到平均值为 0（只看 KEYS 里的 8 个回答，不含外部选项）。"""
    m = mean([scores[k] for k in KEYS])
    return {k: scores[k] - m for k in KEYS}


def rmse(est, truth):
    """
    两组分数的均方根误差 √(平均 (估计 − 真实)²)，两者都先平移到平均值为 0。

    【思路】只有差值有意义，所以比较之前先去掉"整体平移"这个看不见的自由度。
    【返回】float（≥ 0），单位和效用相同。
    """
    e, t = centered(est), centered(truth)
    return math.sqrt(mean([(e[k] - t[k]) ** 2 for k in KEYS]))


# =============================================================================
# 各种"问法"：每个函数模拟一次标注任务，返回这次任务产生的选择事件
# =============================================================================
def task_pair(rng):
    """两两比较：随机挑 2 个回答，问哪个好。成本：读 2 个回答。"""
    a, b = rng.sample(KEYS, 2)
    w, l = sample_pair(TRUE, a, b, rng)
    return [pair_to_event(w, l)]


def task_choose4(rng):
    """多选一：随机挑 4 个回答，挑最好的。成本：读 4 个回答。"""
    opts = rng.sample(KEYS, 4)
    return [choice_to_event(sample_choice(TRUE, opts, rng), opts)]


def task_rank4(rng):
    """排序：随机挑 4 个回答，从好到坏排序（拆成 3 个事件）。成本：读 4 个回答，再加上排序的额外功夫。"""
    opts = rng.sample(KEYS, 4)
    return ranking_to_events(sample_ranking(TRUE, opts, rng))


def task_accept(rng):
    """接受 / 拒绝：随机挑 1 个回答，问能不能接受（外部选项的效用为 THRESHOLD）。成本：读 1 个回答。"""
    k = rng.choice(KEYS)
    utils = dict(TRUE)
    utils[OUTSIDE] = THRESHOLD
    return [accept_reject_to_event(k, sample_accept(utils, k, rng))]


def fit(events, accept=False):
    """
    拟合事件。接受 / 拒绝数据要把外部选项固定为 0（它是零点），其余数据把平均值固定为 0。

    【返回】dict {回答: 分数}（接受 / 拒绝时还包含 ∅: 0）
    """
    if accept:
        return fit_choice_model(events, KEYS, fixed={OUTSIDE: 0.0}, tol=1e-7)
    return fit_choice_model(events, KEYS, tol=1e-7)


# =============================================================================
# 演示 1：五种数据长什么样
# =============================================================================
def demo_formats(rng):
    print_header("演示 1：五种偏好数据，统一写成“从集合中选一个”")
    a, b = "A", "D"
    w, l = sample_pair(TRUE, a, b, rng)
    print(f"  两两比较      {w} ≻ {l}")
    print(f"                → 事件 {pair_to_event(w, l)}")
    opts = ["B", "C", "E", "G"]
    j = sample_choice(TRUE, opts, rng)
    print(f"  多选一        从 {{{', '.join(opts)}}} 里选了 {j}")
    print(f"                → 事件 {choice_to_event(j, opts)}")
    rk = sample_ranking(TRUE, ["A", "C", "F", "H"], rng)
    print(f"  完整排序      {' ≻ '.join(rk)}")
    for e in ranking_to_events(rk):
        print(f"                → 事件 {e}")
    utils = dict(TRUE)
    utils[OUTSIDE] = THRESHOLD
    acc = sample_accept(utils, "F", rng)
    print(f"  接受 / 拒绝    回答 F：{'接受' if acc else '拒绝'}")
    print(f"                → 事件 {accept_reject_to_event('F', acc)}   （∅ 表示“什么都不要”）")
    print("  逐项反馈矩阵  很多用户 × 很多回答，每格是接受(1) / 拒绝(0) / 没看过(空)；每个格子就是一次接受 / 拒绝。")
    print("                 每个用户有自己的口味时，需要“用户 × 回答”的模型（第 4 讲：因子模型）。")
    print("""
观察：
  · 不管数据长什么样，都能拆成“从集合 S 里选了 j”的事件。在 softmax 模型下，事件的概率都是 e^{s_j} / Σ_{k∈S} e^{s_k}，
    所以**一个拟合器**就能处理所有形式（choice_data.fit_choice_model）。
  · 排序被拆成“先挑第一名、再从剩下的挑第二名……”，这就是 Plackett–Luce 模型。
  · 接受 / 拒绝被看成“和外部选项 ∅ 比较”。∅ 的分数固定为 0，其他分数就有了绝对意义（演示 4）。""")


# =============================================================================
# 演示 2：一个拟合器搞定所有形式
# =============================================================================
def demo_one_fitter(rng, n_tasks=3000):
    print_header("演示 2：一个拟合器搞定所有形式（每种形式各做 3000 次任务）")
    results = {}
    results["两两比较"] = fit([e for _ in range(n_tasks) for e in task_pair(rng)])
    results["四选一"] = fit([e for _ in range(n_tasks) for e in task_choose4(rng)])
    rankings = [sample_ranking(TRUE, rng.sample(KEYS, 4), rng) for _ in range(n_tasks)]
    results["四个排序（PL）"] = fit([e for r in rankings for e in ranking_to_events(r)])
    results["排序拆成对（BT）"] = fit([pair_to_event(w, l) for r in rankings for w, l in ranking_to_pairs(r)])
    acc = fit([e for _ in range(n_tasks) for e in task_accept(rng)], accept=True)
    results["接受/拒绝"] = acc

    names = list(results)
    print("所有分数都平移到平均值为 0 后再比较（只有差值有意义）。真实值 = 专家的效用（平移后）。\n")
    print(f"  {'回答':4s} {'真实':>6s} " + " ".join(pad(n, 15, 'right') for n in names))
    for k in KEYS:
        row = " ".join(f"{centered(results[n])[k]:>15.2f}" for n in names)
        print(f"  {k:4s} {TRUE_CENTERED[k]:>6.2f} {row}")
    print(f"  {'误差':4s} {'':>6s} " + " ".join(f"{rmse(results[n], TRUE):>15.3f}" for n in names))
    print("\n  接受 / 拒绝数据的**原始**分数（没有平移，外部选项 ∅ = 0）和“真实效用 − 阈值 3.5”：")
    print("    " + "  ".join(f"{k}: {acc[k]:+.2f}/{TRUE[k] - THRESHOLD:+.2f}" for k in KEYS))
    print("""
观察：
  · 数据足够多时，五种做法都把效用学回来了（误差都只有零点几）。
  · “把排序拆成两两比较、当成 BT 数据拟合”和“按 Plackett–Luce 拟合排序”几乎一样。原因（第 3 讲证明）：
    在“效用 + Gumbel 噪声”的模型下，排序里 a 排在 b 前面的概率**恰好**是 σ(s_a − s_b)。所以拆成对是合理的，
    只是这些对之间不独立（演示 5）。
  · 接受 / 拒绝的原始分数不需要平移，就和“效用 − 阈值”对得上：它学到的是“比及格线高多少”，这是绝对的量。
    但是离阈值最远的 H（真实值 −2.90）估得最不准：它几乎总被拒绝，结果几乎是确定的，信息很少
    （第 7 讲：五五开的问题信息量最大）。""")


# =============================================================================
# 演示 3：同样的成本，哪种问法最好
# =============================================================================
def demo_cost(rng, reps=200, budget=240):
    print_header("演示 3：同样的标注成本，哪种问法学得最好？")
    plans = [
        ("两两比较", task_pair, 2, False),
        ("四选一", task_choose4, 4, False),
        ("四个排序（成本 4）", task_rank4, 4, False),
        ("四个排序（成本 6）", task_rank4, 6, False),
        ("接受/拒绝", task_accept, 1, True),
    ]
    print(f"""成本 = 读了多少个回答（读一个回答花的时间远多于点一下按钮）。预算：读 {budget} 个回答。
  两两比较每次读 2 个 → {budget // 2} 次任务；四选一每次读 4 个 → {budget // 4} 次；
  排序也读 4 个，但排序比挑一个费脑子，所以另算一种"成本 6"（{budget // 6} 次）；接受 / 拒绝每次读 1 个 → {budget} 次。
每种做法重复 {reps} 次，报告平均误差（RMSE，越小越好）和排名一致性（肯德尔 τ，越大越好）。
每次任务的信息量上限：两两 1 比特，四选一 log2(4) = 2 比特，排序 log2(4!) ≈ 4.58 比特，接受 / 拒绝 1 比特。
""")
    print(f"  {pad('问法', 20)} {pad('任务数', 6, 'right')} {pad('RMSE', 8, 'right')} {pad('肯德尔τ', 8, 'right')}")
    for name, task, cost, accept in plans:
        errs, taus = [], []
        n = budget // cost
        for _ in range(reps):
            events = [e for _ in range(n) for e in task(rng)]
            est = fit(events, accept=accept)
            errs.append(rmse(est, TRUE))
            taus.append(kendall_tau({k: est[k] for k in KEYS}, TRUE))
        print(f"  {pad(name, 20)} {n:>6d} {mean(errs):>8.3f} {mean(taus):>8.3f}")
    print("""
观察：
  · 在这个成本模型下，排序最划算：读同样多的回答，排序榨出的信息最多（一次排序 ≈ 4.58 比特）。
    即使排序要多花 50% 的时间（成本 6），它仍然和两两比较差不多甚至更好。这是 InstructGPT 让标注员
    一次排 4~9 个回答的原因之一：读回答是最贵的部分，读完以后多给一点信息很便宜。
  · 四选一只告诉你“谁最好”，扔掉了其余三个之间的顺序，所以不如排序。
  · 接受 / 拒绝每次只读一个，任务数最多，在这个设置里误差反而比两两比较小。但它的信息集中在“离阈值近”的
    回答上：阈值设得离所有回答都很远时（几乎全接受或全拒绝），它就几乎没有信息了。
  · 这些结论**依赖成本模型**：如果排序让标注员疲劳、出错更多（真实的人排 9 个回答时，后面几名常常是乱排的），
    或者读回答很便宜、点按钮很贵，结论就会变。设计数据收集时，最好先做一个这样的小模拟，再做一个小规模的真人试点。""")


# =============================================================================
# 演示 4：接受 / 拒绝能看到"绝对水平"
# =============================================================================
def demo_absolute(rng, n=4000):
    print_header("演示 4：两两比较只告诉你“谁更好”，接受 / 拒绝还告诉你“好不好”")
    worse = {k: v - 3.0 for k, v in TRUE.items()}     # 另一个世界：所有回答都差 3 分（例如换了个更弱的模型）
    for name, world in [("原来的世界", TRUE), ("所有回答都差 3 分的世界", worse)]:
        pairs = []
        accepts = []
        utils = dict(world)
        utils[OUTSIDE] = THRESHOLD
        for _ in range(n):
            a, b = rng.sample(KEYS, 2)
            pairs.append(pair_to_event(*sample_pair(world, a, b, rng)))
            k = rng.choice(KEYS)
            accepts.append(accept_reject_to_event(k, sample_accept(utils, k, rng)))
        est_pair = fit(pairs)
        est_acc = fit(accepts, accept=True)
        rate = mean([1.0 if w != OUTSIDE else 0.0 for w, _ in accepts])
        print(f"\n  【{name}】")
        print("    两两比较拟合（平移后）：" + "  ".join(f"{k}:{est_pair[k]:+.2f}" for k in KEYS))
        print("    接受 / 拒绝拟合（原始）：" + "  ".join(f"{k}:{est_acc[k]:+.2f}" for k in KEYS))
        print(f"    总体接受率：{rate:.1%}")
    print("""
观察：
  · 两个世界的两两比较数据服从**完全相同**的分布（效用整体平移不改变任何胜率），拟合结果也几乎一样。
    只看两两比较，你永远发现不了“这批回答整体都很差”。
  · 接受 / 拒绝数据一下子就看出来了：第二个世界里几乎所有回答都被拒绝，分数整体低了约 3。
  · 实践中常把两者结合：Chatbot Arena 的投票除了“A 好 / B 好 / 平局”，还有“两个都不好”；
    很多产品同时收集点赞 / 点踩和两两比较。第 4 讲会看到，接受 / 拒绝的“逐项数据”还能估计每个用户的“胃口”。""")


# =============================================================================
# 演示 5：排序拆出来的对不独立
# =============================================================================
def bootstrap_sd(units, to_pairs, rng, rounds=200):
    """
    自助法估计 s_A − s_B 的标准差：对"数据单元"有放回地重抽样，每次把单元变成两两比较、用 BT 拟合。

    【参数】
        units   ：数据单元的列表（例如一次排序就是一个单元，或者一条比较就是一个单元）
        to_pairs：把一个单元变成若干 (赢家, 输家) 的函数
        rng、rounds：随机数生成器和重抽样次数
    【返回】float：s_A − s_B 的估计值在重抽样中的标准差
    【边界】自助法必须对**真正独立的单元**重抽样。同一次排序产生的 6 对比较不独立，应该整体一起抽。
    """
    vals = []
    for _ in range(rounds):
        sample = rng.choices(units, k=len(units))
        est = fit([pair_to_event(w, l) for u in sample for w, l in to_pairs(u)])
        vals.append(est["A"] - est["B"])
    m = mean(vals)
    return math.sqrt(mean([(v - m) ** 2 for v in vals]))


def demo_dependence(rng, n_rankings=150):
    print_header("演示 5：排序拆出来的两两比较不是独立的")
    rankings = [sample_ranking(TRUE, rng.sample(KEYS, 4), rng) for _ in range(n_rankings)]
    pairs = [p for r in rankings for p in ranking_to_pairs(r)]
    right = bootstrap_sd(rankings, ranking_to_pairs, rng)
    naive = bootstrap_sd(pairs, lambda p: [p], rng)
    print(f"{n_rankings} 次“四个回答排序”，拆成 {len(pairs)} 对比较，用 BT 拟合。用自助法估计 s_A − s_B 的标准差：")
    print(f"  按“排序”整体重抽样（正确）：      {right:.3f}")
    print(f"  按“对”重抽样（当成独立的比较）：  {naive:.3f}   （是正确值的 {naive / right:.0%}）")
    print("""
观察：
  · 同一次排序产生的 6 对比较共享同一次“感受”（同一组噪声），它们是相关的。把它们当成 6 次独立的比较，
    等于假装自己有更多的数据，估计出来的不确定性就偏小（这里低估了约十分之一）—— 置信区间太窄，过于自信。
    每次排的回答越多，同一次排序拆出的对越多、相关性越强，低估得也越厉害（编写讲义时另做的一次运行：
    每次排 2 个时两种算法一样（这时本来就没有相关性），排 4 个时约 86%，排 8 个时约 78%）。
  · InstructGPT 的论文报告了一个相关的现象：把同一个提示词的所有比较打散混进数据集，奖励模型很快就过拟合了；
    他们改成把同一个提示词的 K(K−1)/2 对比较放在同一个批次里一起训练。
  · 原则：自助法、交叉验证、训练 / 测试划分，都要以**真正独立的单元**为单位（这里是“一次排序”，
    在 RLHF 里通常是“一个提示词”）。第 5 讲做交叉验证时还会遇到这个原则。""")


def main():
    rng = random.Random(1952)
    demo_formats(rng)
    demo_one_fitter(rng)
    demo_cost(rng)
    demo_absolute(rng)
    demo_dependence(rng)


if __name__ == "__main__":
    main()