"""
=============================================================================
第 9 讲 · 代码 5：社会选择的视角 —— DPO / BT 隐式地用 Borda 计数聚合；偏好成环时的 Nash 策略
=============================================================================

对应讲义：lectures/lecture09_dpo.md 第 7 节（预告第 16、17 讲）

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
偏好数据来自很多人，人和人的偏好不一样（第 1、4 讲）。如果标注数据里**没有记录**"这条是谁标的"
（"隐藏的情境"），用 BT / DPO 拟合，学到的"奖励"是在用什么规则把不同人的偏好合起来？

    演示 1：两类用户、四个候选回答。BT 学到的排序 = Borda 计数的排序（Siththaranjan 等 2024），
            它可能和"平均效用"的排序不同。随机生成 2000 个例子检查这个结论。
    演示 2：Borda / BT 只看"多少人更喜欢"，不看"有多喜欢"：少数人的强烈偏好会被多数人的微弱偏好压过。
    演示 3：偏好成环（A 胜 B、B 胜 C、C 胜 A）时，BT 给出一个"冠军"，可它会被别的回答打败；
            Nash 均衡（"最大彩票"）是一个混合策略，对任何回答的胜率都不低于一半（Nash 偏好学习，Munos 等 2024）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_social_choice.py        # 约 2 秒

只需要 numpy，不依赖其他文件。
"""

import unicodedata

import numpy as np

NAMES = ["A", "B", "C", "D"]


def print_header(title):
    """打印一个带分隔线的标题（和 toy_lm.py 里的一样；这里再写一份，让本文件只依赖 numpy）。"""
    print()
    print("=" * 78)
    print(f"  {title}")
    print("=" * 78)


def pad(text, width, align="left"):
    """按显示宽度补空格（汉字算 2 个宽度）。"""
    w = sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text)
    spaces = " " * max(0, width - w)
    return text + spaces if align == "left" else spaces + text


def sigmoid(x):
    """σ(x)。"""
    return 1.0 / (1.0 + np.exp(-x))


def pairwise_matrix(groups):
    """
    隐藏情境下的总体偏好矩阵：P[i, j] = Σ_g w_g · σ(u_g[i] − u_g[j])（对角线为 0）。

    【参数】groups：[(人群比例 w_g, 效用数组 u_g), ...]
    【思路】标注员从各人群里按比例随机抽，但数据里不记录是哪个人群 —— 我们只能看到混合之后的胜率。
    """
    M = len(groups[0][1])
    P = np.zeros((M, M))
    for w, u in groups:
        P += w * sigmoid(u[:, None] - u[None, :])
    np.fill_diagonal(P, 0.0)
    return P


def bt_fit(P, iters=30):
    """
    在"所有对都被等量比较"的无限数据上拟合 BT：最大化 Σ_{i≠j} P[i,j] · log σ(s_i − s_j)。牛顿法。

    【返回】分数（平均值为 0）
    【偏好学习联系】DPO 在数据无限、模型足够灵活时，隐式奖励 β·log(π/π_ref) 就是这个 BT 分数（差一个常数）。
    """
    M = len(P)
    s = np.zeros(M)
    for _ in range(iters):
        p = sigmoid(s[:, None] - s[None, :])
        np.fill_diagonal(p, 0.0)
        g = (P - p).sum(axis=1)
        W = p * p.T                                   # p_ij (1 − p_ij) = p_ij p_ji
        H = W - np.diag(W.sum(axis=1))
        s = s - np.linalg.pinv(H) @ g
    return s - s.mean()


def borda(P):
    """Borda 计数（用胜率）：每个候选的得分 = 它对其他每个候选的胜率之和。"""
    return P.sum(axis=1)


# =============================================================================
# 演示 1
# =============================================================================
def demo_borda():
    print_header("演示 1：隐藏的情境 —— BT 学到的排序就是 Borda 计数的排序")
    groups = [(0.57, np.array([0.4, 0.4, 0.7, -1.2])), (0.43, np.array([-2.0, -1.3, -3.5, 5.2]))]
    print("两类用户（数据里不记录谁是谁）：")
    print(f"  · 第一类（57%）：效用 {groups[0][1].tolist()}（偏好都不强烈，最不喜欢 D）")
    print(f"  · 第二类（43%）：效用 {groups[1][1].tolist()}（非常喜欢 D）\n")
    P = pairwise_matrix(groups)
    s = bt_fit(P)
    b = borda(P)
    mean_u = sum(w * u for w, u in groups)
    print("  总体胜率矩阵 P[行 胜 列]：")
    print("        " + "".join(f"{n:>8s}" for n in NAMES))
    for i, n in enumerate(NAMES):
        print(f"    {n:>3s} " + "".join(f"{P[i, j]:>8.3f}" if i != j else f"{'—':>8s}" for j in range(4)))
    pi = np.exp(s) / np.exp(s).sum()                 # β = 1、π_ref 均匀时 DPO 的最优策略
    print(f"\n  {pad('候选', 6)} {pad('BT 分数', 9, 'right')} {pad('Borda 得分', 11, 'right')} {pad('平均效用', 9, 'right')} "
          f"{pad('DPO 最优策略（β=1，π_ref 均匀）', 32, 'right')}")
    for i, n in enumerate(NAMES):
        print(f"  {pad(n, 6)} {s[i]:>+9.3f} {b[i]:>11.3f} {mean_u[i]:>+9.3f} {pi[i]:>32.3f}")
    order = lambda x: " > ".join(NAMES[k] for k in np.argsort(-x))    # noqa: E731
    print(f"\n  排序：BT {order(s)}；Borda {order(b)}；平均效用 {order(mean_u)}")

    # 随机检查
    rng = np.random.default_rng(0)
    same, exceptions = 0, []
    for _ in range(2000):
        w = rng.uniform(0.5, 0.7)
        u1 = rng.normal(0, 1, 4)
        u2 = rng.normal(0, 3, 4)
        Pr = pairwise_matrix([(w, u1), (1 - w, u2)])
        sr, br = bt_fit(Pr), borda(Pr)
        if np.array_equal(np.argsort(-sr), np.argsort(-br)):
            same += 1
        else:
            gaps = np.diff(np.sort(br))
            exceptions.append(float(gaps.min()))
    print(f"\n  随机生成 2000 个“两类用户、四个候选”的例子：BT 排序和 Borda 排序完全相同的有 {same} 个。")
    if exceptions:
        print(f"  不相同的 {len(exceptions)} 个例子里，Borda 得分最接近的两个候选只差 "
              + "、".join(f"{g:.4f}" for g in exceptions) + "（几乎并列）。")
    print("""
观察：
  · BT（因此也是 DPO）在混合了不同人群的数据上学到的排序，就是 Borda 计数的排序：每个候选的得分
    是“它赢其他候选的概率之和”。Siththaranjan、Laidlaw、Hadfield-Menell（2024）把这叫做
    “隐藏情境下的偏好学习隐式地用 Borda 计数聚合”。随机检查的 2000 个例子全部一致。
  · 为什么？这就是第 5 讲的“矩匹配”：最大似然处，每个候选的预期胜场 = 实际胜场。所有对都被等量比较时，
    “实际胜场”正是 Borda 得分，而“预期胜场”Σ_j σ(s_i − s_j) 随 s_i 单调递增 —— 所以两者的排序必然相同。
  · 平均效用（“功利主义”的聚合）的排序可以不同：第二类用户对 D 的强烈喜爱让 D 的平均效用最高，
    但 Borda / BT 把 B 排第一，D 第二 —— 第一类用户“温和地不喜欢 D”的人数更多。
  · 哪个是“对的”？这是一个**价值选择**，不是技术问题（第 16、17 讲）。重要的是知道：
    用 BT / DPO 拟合混合人群的数据，你已经（不知不觉地）选择了 Borda 计数。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_intensity():
    print_header("演示 2：只看“多少人更喜欢”，不看“有多喜欢”")
    groups = [(0.6, np.array([3.0, 0.0])), (0.4, np.array([0.0, 30.0]))]
    P = pairwise_matrix(groups)
    s = bt_fit(P)
    mean_u = sum(w * u for w, u in groups)
    print("两个回答 A、B。60% 的用户稍微更喜欢 A（效用差 3），40% 的用户极其喜欢 B（效用差 30）。\n")
    print(f"  P(A 胜 B) = 0.6 × σ(3) + 0.4 × σ(−30) = {P[0, 1]:.3f}")
    print(f"  BT 分数差 s_A − s_B = {s[0] - s[1]:+.3f}（BT / DPO 更偏向 A）")
    print(f"  平均效用差 = {mean_u[0] - mean_u[1]:+.1f}（功利主义更偏向 B）")
    print("""
观察：
  · 成对比较只记录“谁赢了”，不记录“赢了多少”。σ 把再强烈的偏好也压到 1 以内：效用差 3 和效用差 30，
    在胜率上只是 0.95 和 1.00 的区别。所以多数人的微弱偏好可以压过少数人的强烈偏好。
  · 这和投票一样：一人一票，不论多在乎。如果你想体现“强度”，需要别的数据（打分、愿意付出的代价……，第 2 讲）
    或者显式地对不同人群建模（第 4 讲的个性化、分布式偏好学习）。""")


# =============================================================================
# 演示 3
# =============================================================================
def demo_cycle():
    print_header("演示 3：偏好成环 —— BT 的冠军会被打败；Nash 混合策略不会")
    # P(A≻B) = 0.9，P(B≻C) = 0.6，P(C≻A) = 0.6：一个不对称的环
    P = np.array([[0.0, 0.9, 0.4],
                  [0.1, 0.0, 0.6],
                  [0.6, 0.4, 0.0]])
    names = NAMES[:3]
    s = bt_fit(P)
    print("三个回答的总体胜率：P(A 胜 B) = 0.9，P(B 胜 C) = 0.6，P(C 胜 A) = 0.6 —— A 胜 B、B 胜 C、C 胜 A，成环。")
    print(f"（这样的环可以由三类用户的不同偏好混合产生：第 1 讲的孔多塞悖论。）\n")
    print("  BT 分数：" + "，".join(f"{n} {v:+.3f}" for n, v in zip(names, s)))
    # 最大彩票 / Nash 均衡：对反对称矩阵 M = P − Pᵀ（偏好的"净优势"），找 p ≥ 0 使 pᵀ M ≥ 0。
    # 三个候选成环时有封闭解：p ∝ (M_BC, M_CA, M_AB)（练习 7）
    M = P - P.T
    np.fill_diagonal(M, 0.0)
    nash = np.array([M[1, 2], M[2, 0], M[0, 1]])
    nash = nash / nash.sum()

    def win_rate(p, q):
        """策略 p（混合）对策略 q 的期望胜率：Σ_ij p_i q_j P[i,j]，自己对自己算 0.5。"""
        PP = P + 0.5 * np.eye(len(P))
        return float(p @ PP @ q)

    pures = [np.eye(3)[k] for k in range(3)]
    candidates = [("只选 A（BT 的冠军）", pures[0]), ("只选 C", pures[2]),
                  ("BT 的 softmax（β = 1）", np.exp(s) / np.exp(s).sum()),
                  ("BT 的 softmax（β = 0.2）", np.exp(s / 0.2) / np.exp(s / 0.2).sum()),
                  ("Nash 混合策略", nash)]
    print(f"\n  {pad('策略', 24)} {pad('选 A / B / C 的概率', 22, 'right')} " + "".join(pad(f"对 {n} 的胜率", 13, "right") for n in names)
          + pad("最差情况", 10, "right"))
    for name, p in candidates:
        rates = [win_rate(p, q) for q in pures]
        probs = " / ".join(f"{v:.2f}" for v in p)
        print(f"  {pad(name, 24)} {probs:>22s} " + "".join(f"{r:>13.3f}" for r in rates) + f"{min(rates):>10.3f}")
    print("""
观察：
  · BT 必须给每个回答一个数，所以它把环“摊平”了：A 的分数最高（它大胜 B），DPO 会把概率往 A 推。
    可是 A 对 C 只有 40% 的胜率 —— 一个总是选 A 的策略会被“总是选 C”打败。
  · BT 的 softmax 越尖锐（β 越小，越接近只选 A），最差情况越糟。
  · Nash 混合策略（社会选择里叫“最大彩票”，Fishburn 1984）：A 1/6、B 1/6、C 2/3。它对任何回答的胜率都至少是一半 ——
    没有任何回答（也没有任何混合策略）能在平均意义上打败它。直觉：C 是 A 的克星，而 A 是最“强”的，
    所以要多放 C；B 被 A 大胜，C 又被 B 小胜……均衡在这些相互克制之间取得平衡。
  · Nash 偏好学习（Munos 等 2024）、自博弈偏好优化（SPPO 等）就是用这个思想训练语言模型：
    不假设 BT（不假设偏好可以用一个分数表示），而是直接找“不会被任何别的策略打败”的策略。
    当偏好确实可传递时，它退化成“选孔多塞赢家”；不可传递时，它给出一个有原则的混合。""")


def main():
    demo_borda()
    demo_intensity()
    demo_cycle()


if __name__ == "__main__":
    main()