"""
=============================================================================
第 4 讲 · 代码 5：成对比较的个性化 —— 一个 BT 给所有人，还是每人一个口味向量？
=============================================================================

对应讲义：lectures/lecture04_factor_models.md 第 6 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
RLHF 的数据是两两比较。在因子模型下，用户 i 比较回答 j 和 k：

        P(j ≻ k | i) = σ( a_i + z_j + U_i·V_j − (a_i + z_k + U_i·V_k) )
                     = σ( (z_j − z_k) + U_i·(V_j − V_k) )

用户的胃口 a_i 消掉了（Rasch 的"特定客观性"），但口味 U_i 没有消掉 —— 不同的用户会把同一对回答排成不同的顺序。
教材第 1.10 节总结：只有当 (1) 模型是 Rasch（没有口味项）、(2) 只有一个用户、或 (3) 所有用户的口味相同时，
用户参数才会完全消掉，标准的 BT 才是对的。

    演示 1：胃口相同、口味相同的用户，两两比较的数据完全一样；胃口不同、口味相同 —— 也完全一样。
            口味不同 —— 就不一样了。
    演示 2：三种模型比较 —— 一个 BT 给所有人（今天的奖励模型）、每人单独一个 BT、
            个性化的因子 BT（每人一个口味向量，回答的向量大家共享）。每人的比较次数从 5 到 80。
    演示 3：谁被"一个 BT 给所有人"亏待了？口味离"平均人"越远的用户，全局 BT 对他们越不准。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_pairwise_personalization.py        # 约 3 秒

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

import math
import random

from factor_models import TasteWorld, dot, fit_pairwise_mf, pairwise_prob
from toy_world import fit_bradley_terry, mean, pad, print_header, sigmoid


def compare(world, i, j, k, rng):
    """用户 i 比较回答 j 和 k，返回 (赢家, 输家)。胃口 a_i 在差值里消掉了，所以这里不用它。"""
    p = sigmoid(world.logit(i, j) - world.logit(i, k))
    return (j, k) if rng.random() < p else (k, j)


def make_comparisons(world, m, rng):
    """每个用户做 m 次比较（随机挑两个不同的回答）。返回 list[(用户, 赢家, 输家)]。"""
    out = []
    for i in world.users:
        for _ in range(m):
            j, k = rng.sample(world.items, 2)
            w, l = compare(world, i, j, k, rng)
            out.append((i, w, l))
    return out


def test_pairs(world, rng, n_per_user=30):
    """每个用户随机抽 n 对回答作为测试（只用来判断"模型认为谁更好"是否和用户的真实效用一致）。"""
    return [(i, *rng.sample(world.items, 2)) for i in world.users for _ in range(n_per_user)]


def accuracy_for(pairs, score_fn, world):
    """
    排序正确率：在测试对上，模型认为更好的那个，是否真的是这个用户效用更高的那个。

    【参数】pairs：list[(用户, j, k)]；score_fn(i, j, k) → 模型认为 "j ≻ k" 的概率；world：真实世界
    【返回】0~1
    """
    hits = 0.0
    for i, j, k in pairs:
        p = score_fn(i, j, k)
        truth = world.logit(i, j) > world.logit(i, k)
        hits += 0.5 if p == 0.5 else (1.0 if (p > 0.5) == truth else 0.0)
    return hits / len(pairs)


# =============================================================================
# 演示 1：胃口消掉，口味消不掉
# =============================================================================
def demo_cancellation(world):
    print_header("演示 1：两两比较里，胃口消掉了，口味没有")
    j, k = world.items[0], world.items[1]
    base = world.users[0]
    taste = world.taste[base]
    print(f"三个虚构的用户比较回答 {j} 和 {k}：")
    rows = [("用户甲", 2.0, taste), ("用户乙（更挑剔，口味相同）", -2.0, taste),
            ("用户丙（胃口同甲，口味相反）", 2.0, [-t for t in taste])]
    for name, a, t in rows:
        lj = a + 2.0 * world.quality[j] + dot(t, world.style[j])
        lk = a + 2.0 * world.quality[k] + dot(t, world.style[k])
        print(f"  {pad(name, 30)} 胃口 {a:+.1f}  接受 {j} 的概率 {sigmoid(lj):.3f}  接受 {k} 的概率 {sigmoid(lk):.3f}"
              f"  P({j} ≻ {k}) = {sigmoid(lj - lk):.3f}")
    print("""
观察：
  · 甲和乙的胃口差了 4，接受的概率天差地别，可是两两比较的概率完全一样：胃口在“j 的 logit − k 的 logit”里消掉了。
    所以两两比较的数据**学不到**用户的胃口（宽严），也不受它的干扰。
  · 丙的口味和甲相反，两两比较的结果也反过来了。口味（U_i）在差值里**没有**消掉：
    (z_j − z_k) + U_i·(V_j − V_k) 依赖于 U_i。这就是为什么“一个 BT 给所有人”在口味有分歧时会出问题。""")


# =============================================================================
# 演示 2：三种模型
# =============================================================================
def demo_models(world, rng):
    print_header("演示 2：一个 BT 给所有人 vs 每人一个 BT vs 个性化因子 BT")
    tests = test_pairs(world, rng)
    print(f"{len(world.users)} 个用户、{len(world.items)} 个回答。每个用户做 m 次两两比较；测试时每人 30 对，看排序正确率。\n")
    print(f"  {pad('每人比较 m 次', 14)} {pad('全局 BT', 9, 'right')} {pad('每人单独 BT', 12, 'right')} {pad('个性化因子 BT（K=2）', 21, 'right')}")
    results = {}
    for m in (5, 20, 80):
        comps = make_comparisons(world, m, rng)
        glob = fit_bradley_terry([(w, l) for _, w, l in comps], world.items, l2=1e-3)
        per_user = {}
        for i in world.users:
            mine = [(w, l) for u, w, l in comps if u == i]
            per_user[i] = fit_bradley_terry(mine, world.items, l2=0.05, max_iters=3000, tol=1e-6)
        fac = fit_pairwise_mf(comps, world.users, world.items, 2, 1.0, random.Random(5), sweeps=20)
        acc_g = accuracy_for(tests, lambda i, j, k: sigmoid(glob[j] - glob[k]), world)
        acc_u = accuracy_for(tests, lambda i, j, k: sigmoid(per_user[i][j] - per_user[i][k]), world)
        acc_f = accuracy_for(tests, lambda i, j, k: pairwise_prob(fac, i, j, k), world)
        results[m] = (glob, fac)
        print(f"  {m:>14d} {acc_g:>9.1%} {acc_u:>12.1%} {acc_f:>21.1%}")
    print("""
观察：
  · 全局 BT（今天绝大多数奖励模型的做法）：数据再多，也到不了很高 —— 它只能学到“平均人”的偏好，
    对口味和平均人不同的用户，排序注定有一部分是反的。
  · 每人单独一个 BT：每人 40 个回答要估计 40 个分数，几次比较根本不够；数据很多时才慢慢变好。
  · 个性化因子 BT：回答的向量 V_j 大家共享（用所有人的数据估计），每个用户只需要估计 2 维的口味 U_i。
    每人只比较 5 次时，它反而不如全局 BT：5 次比较估不准 2 维口味，估出来的口味主要是噪声，
    还不如干脆假设“这个人和平均人一样”。从 20 次起它就明显超过全局 BT，这是“从别人那里借信息”的力量。
    实践中常在两者之间折中：数据少的用户，口味向量被正则化拉向 0（也就是拉向平均人），数据多了再放开。""")
    return tests, results


# =============================================================================
# 演示 3：谁被亏待了
# =============================================================================
def demo_who_loses(world, tests, results):
    print_header("演示 3：谁被“一个 BT 给所有人”亏待了？")
    glob, fac = results[80]
    avg_taste = [mean([world.taste[i][d] for i in world.users]) for d in range(2)]

    def dist(i):
        return math.sqrt(sum((world.taste[i][d] - avg_taste[d]) ** 2 for d in range(2)))

    ranked = sorted(world.users, key=dist)
    third = len(ranked) // 3
    print("把用户按“口味离平均口味有多远”分成三组（每人比较 80 次的模型）：\n")
    print(f"  {pad('用户组', 18)} {pad('平均距离', 8, 'right')} {pad('全局 BT', 9, 'right')} {pad('个性化因子 BT', 14, 'right')}")
    for name, grp in [("口味最接近平均的 1/3", ranked[:third]), ("中间的 1/3", ranked[third:2 * third]),
                      ("口味最“小众”的 1/3", ranked[2 * third:])]:
        gs = set(grp)
        rows = [t for t in tests if t[0] in gs]
        a_g = accuracy_for(rows, lambda i, j, k: sigmoid(glob[j] - glob[k]), world)
        a_f = accuracy_for(rows, lambda i, j, k: pairwise_prob(fac, i, j, k), world)
        print(f"  {pad(name, 18)} {mean([dist(i) for i in grp]):>8.2f} {a_g:>9.1%} {a_f:>14.1%}")
    print("""
观察：
  · 全局 BT 对口味接近平均的用户相当准，对口味“小众”的用户差得多 —— 平均准确率掩盖了这种不平等。
  · 个性化模型对三组用户几乎一样准，小众用户甚至略高：他们的口味很鲜明（口味向量长），比较结果更“一边倒”，
    80 次比较足够把它学出来。个性化消除了全局模型造成的不平等 —— 前提是每个人都有足够的数据（演示 2）。
  · 这把第 1 讲的“谁的偏好”问题变成了可以测量的东西：报告奖励模型的准确率时，按用户群体分开报告
    （第 17 讲：公平）。个性化是一种回答；另一种回答是承认“一个模型服务所有人”本身是一个价值选择（第 16 讲）。""")


def main():
    rng = random.Random(2024)
    world = TasteWorld(rng)
    demo_cancellation(world)
    tests, results = demo_models(world, rng)
    demo_who_loses(world, tests, results)


if __name__ == "__main__":
    main()