"""
=============================================================================
第 2 讲 · 代码 4：Bradley–Terry 模型的三个细节 —— 比较图、拟合优度、平局
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
有了比较数据，"拟合一个 BT 模型"看起来只是一行代码。但有三件事必须先想清楚：

    演示 1：比较图 —— 哪些选项之间比较过？如果某个选项从来没输过，或者选项分成了互不比较的两群，
            "最好的分数"根本不存在，或者不唯一。判断方法是 Zermelo（1929）/ Ford（1957）的连通性条件。
    演示 2：拟合优度 —— 数据真的服从 BT 吗？把 BT 和"每一对各有自己的胜率"的饱和模型比较
            （似然比检验），并检查随机传递性。两群口味不同的人混在一起时，检验会报警。
    演示 3：平局 —— 标注员说"差不多"怎么办？丢掉、各算半场、还是用 Rao–Kupper / Davidson 模型？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_bt_details.py        # 约 5 秒

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

import itertools
import math
import random
from collections import Counter

from choice_data import fit_choice_model, gradient_ascent, pair_to_event
from toy_world import BY_KEY, KEYS, Annotator, make_expert, mean, pad, print_header, sigmoid

EXPERT = make_expert()
TRUE = {k: EXPERT.utility(BY_KEY[k]) for k in KEYS}


def sample_bt(utils, a, b, rng):
    """按 BT 模型模拟一次比较，返回 (赢家, 输家)。"""
    return (a, b) if rng.random() < sigmoid(utils[a] - utils[b]) else (b, a)


# =============================================================================
# 演示 1：比较图
# =============================================================================
def reachable(start, edges):
    """
    从 start 出发，沿着有向边能走到的所有节点（广度优先搜索）。

    【参数】start：起点；edges：dict {节点: 它指向的节点集合}
    【返回】set
    【Python 小知识】广度优先搜索（BFS）：用一个"待访问"列表，每次取出一个节点，把它没访问过的邻居加进去。
    """
    seen = {start}
    todo = [start]
    while todo:
        x = todo.pop()
        for y in sorted(edges.get(x, ())):
            if y not in seen:
                seen.add(y)
                todo.append(y)
    return seen


def check_graph(comparisons, keys):
    """
    检查比较图的两种连通性。

    【思路】
        · 无向连通：把"比较过"看成一条无向边。如果图分成了几块，块与块之间的相对高低完全无从得知。
        · 强连通（Zermelo–Ford 条件）：把"a 赢过 b"看成一条 a → b 的有向边。BT 的最大似然估计存在且唯一
          （在"平均值为 0"的约束下），当且仅当这个有向图是强连通的：从任何一个选项出发，都能沿着
          "赢过"的箭头走到任何其他选项。等价的说法：不管怎样把选项分成两组，总有第一组的某个选项
          赢过第二组的某个选项，**并且反过来也有**。
          如果某个选项从来没输过，就没有箭头指向它 → 不是强连通 → 它的分数会被推到 +∞。
    【参数】comparisons：list[(赢家, 输家)]；keys：全部选项
    【返回】(无向连通: bool, 强连通: bool)
    """
    und = {k: set() for k in keys}
    win = {k: set() for k in keys}
    for w, l in comparisons:
        und[w].add(l)
        und[l].add(w)
        win[w].add(l)
    k0 = keys[0]
    connected = reachable(k0, und) == set(keys)
    strong = all(reachable(k, win) == set(keys) for k in keys)
    return connected, strong


def demo_graph(rng):
    print_header("演示 1：比较图 —— 什么时候“最好的分数”存在？")
    four = ["A", "B", "C", "D"]

    # (a) A 从来没输过：只有专家，而且 A 的效用远高于其他
    utils = {"A": 9.0, "B": 2.0, "C": 1.0, "D": 0.0}
    data = [sample_bt(utils, *rng.sample(four, 2), rng) for _ in range(300)]
    data = [(w, l) for w, l in data if l != "A"]          # 保证 A 一次都没输（为了演示）
    conn, strong = check_graph(data, four)
    print(f"(a) A 一次都没输过（{len(data)} 条比较）。无向连通：{conn}；强连通：{strong}")
    events = [pair_to_event(w, l) for w, l in data]
    for iters in [100, 1000, 10000]:
        s = fit_choice_model(events, four, l2=0.0, max_iters=iters, tol=0.0)
        print(f"    不加正则化，迭代 {iters:>5d} 步：s_A − s_B = {s['A'] - s['B']:6.2f}")
    s = fit_choice_model(events, four, l2=0.01)
    print(f"    加 L2 正则化（0.01）：             s_A − s_B = {s['A'] - s['B']:6.2f}（收敛到一个有限值）")

    # (b) 两群互不比较
    group1, group2 = ["A", "B", "C"], ["F", "G", "H"]
    data2 = [sample_bt(TRUE, *rng.sample(group1, 2), rng) for _ in range(400)]
    data2 += [sample_bt(TRUE, *rng.sample(group2, 2), rng) for _ in range(400)]
    six = group1 + group2
    conn, strong = check_graph(data2, six)
    s = fit_choice_model([pair_to_event(w, l) for w, l in data2], six, l2=0.001)
    print(f"\n(b) 两群互不比较：{group1} 内部比较 400 次，{group2} 内部比较 400 次。无向连通：{conn}；强连通：{strong}")
    print(f"    真实：A 胜 H 的概率 = {sigmoid(TRUE['A'] - TRUE['H']):.3f}；模型的预测 = {sigmoid(s['A'] - s['H']):.3f}")
    for n_bridge in [30, 300]:
        bridge = [sample_bt(TRUE, "C", "F", rng) for _ in range(n_bridge)]
        s2 = fit_choice_model([pair_to_event(w, l) for w, l in data2 + bridge], six, l2=0.001)
        conn, strong = check_graph(data2 + bridge, six)
        print(f"    加上 {n_bridge:>3d} 次 C 与 F 的“桥接”比较之后：强连通：{strong}；模型的预测 = {sigmoid(s2['A'] - s2['H']):.3f}")
    print("""
观察：
  · (a) 一个从来没输过的选项，最大似然会说“它比别人好无穷多”：迭代越久，分数差越大，永远不会停。
    这不是代码的错误，而是数据真的支持“越大越好”。加一点 L2 正则化（相当于“分数不会无穷大”的先验，第 5 讲），
    分数就收敛了。小数据集上（比如一个新模型只打了几场全胜），这种情况非常常见。
  · (b) 两群选项内部比较得再多，群与群之间的高低也完全无从得知。正则化会把两群的平均分都拉向 0，
    于是模型对“A 胜 H”远没有真实的 99% 那么有把握。加上桥接比较以后，群与群之间的差距就能估计了；
    桥接越多越准：30 次就把预测拉到了九成以上，300 次接近真实值。注意 C 和 F 的实力差得很远（C 几乎总赢），
    这样的比较信息量很少（第 7 讲：五五开的比较信息量最大），换成实力接近的两个选项做桥接会更省。
  · 实践建议：拟合之前先检查比较图（check_graph）；设计数据收集时保证图连通，最好让每个选项都和
    各种水平的对手比较过（第 7 讲会讲怎样挑比较最省钱）。""")


# =============================================================================
# 演示 2：拟合优度
# =============================================================================
def saturated_log_lik(counts):
    """
    "饱和模型"的对数似然：每一对 (a, b) 都有自己独立的胜率，取观察到的频率（这是它的最大似然）。

    【参数】counts：dict {(a, b): a 赢 b 的次数}（a、b 两个方向都要有）
    【返回】float：总对数似然（不是平均）
    【边界】某个方向次数为 0 时，0·log(0) 按 0 处理。
    """
    total = 0.0
    for a, b in itertools.combinations(sorted({x for x, _ in counts}), 2):
        nab, nba = counts.get((a, b), 0), counts.get((b, a), 0)
        n = nab + nba
        for c in (nab, nba):
            if c > 0:
                total += c * math.log(c / n)
    return total


def goodness_of_fit(comparisons, keys):
    """
    似然比检验（G 检验）：BT 模型 vs 饱和模型。

    【思路】饱和模型有 M(M−1)/2 个参数（每一对一个胜率），BT 只有 M − 1 个（分数差）。如果数据真的服从 BT，
        饱和模型多出来的参数只能拟合噪声，两者的对数似然差得不多。统计量 G² = 2 × (LL_饱和 − LL_BT)
        在"数据服从 BT"时近似服从自由度为 M(M−1)/2 − (M−1) 的卡方分布。
        M = 8 时自由度是 28 − 7 = 21，卡方分布在 5% 显著性水平下的临界值约为 32.7：G² 远大于它，
        就说明 BT 拟合得不好。
    【参数】comparisons：list[(赢家, 输家)]；keys：全部选项
    【返回】(G², 自由度)
    【边界】需要每一对都比较过足够多次，卡方近似才可靠。
    """
    counts = Counter(comparisons)
    s = fit_choice_model([pair_to_event(w, l) for w, l in comparisons], keys, l2=0.0, tol=1e-9)
    ll_bt = math.fsum(math.log(sigmoid(s[w] - s[l])) for w, l in comparisons)
    m = len(keys)
    return 2 * (saturated_log_lik(counts) - ll_bt), m * (m - 1) // 2 - (m - 1)


def sst_violations(comparisons, keys, z=2.0):
    """
    在数据上数"强随机传递性"的明显违反：P̂(x≻y) ≥ 0.5、P̂(y≻z) ≥ 0.5，但 P̂(x≻z) 比两者中较大的那个
    还小 z 个标准误以上。

    【参数】comparisons；keys；z：几个标准误才算"明显"
    【返回】(违反的三元组个数, 被检查的三元组个数)
    【边界】用观察频率估计胜率，每一对的比较次数太少时标准误很大，几乎检测不到违反。
    """
    counts = Counter(comparisons)
    def phat(a, b):
        n = counts[(a, b)] + counts[(b, a)]
        return counts[(a, b)] / n, n
    bad = checked = 0
    for x, y, w in itertools.permutations(keys, 3):
        pxy, _ = phat(x, y)
        pyw, _ = phat(y, w)
        pxw, n = phat(x, w)
        if pxy >= 0.5 and pyw >= 0.5:
            checked += 1
            se = math.sqrt(max(pxw * (1 - pxw), 1e-6) / n)
            if pxw < max(pxy, pyw) - z * se:
                bad += 1
    return bad, checked


def demo_goodness(rng, per_pair=300):
    print_header("演示 2：数据真的服从 BT 吗？—— 似然比检验与随机传递性")
    concise = Annotator(name="简洁派", taste=-4.0)
    detailed = Annotator(name="详尽派", taste=+4.0)
    # 每个"世界"是一个函数：给一对回答，返回一次比较的结果 (赢家, 输家)。
    # 【Python 小知识】lambda a, b: ... 定义一个没有名字的小函数；字典的值可以是函数。
    worlds = {
        "专家（数据确实来自 BT）": lambda a, b: sample_bt(TRUE, a, b, rng),
        # 混合人群：每次比较先随机抽一个人（60% 简洁派），再按他心里的效用做 BT 比较
        "两群人混合（60% 简洁 + 40% 详尽）": lambda a, b: sample_bt(_mixture_utils(rng, concise, detailed), a, b, rng),
    }
    print(f"每一对回答各比较 {per_pair} 次（8 个回答，28 对，共 {28 * per_pair} 次）。\n")
    print(f"  {pad('数据来源', 34)} {pad('G²', 8, 'right')} {pad('自由度', 6, 'right')} {pad('结论（临界值 32.7）', 20, 'right')} "
          f"{pad('明显违反 SST', 12, 'right')}")
    for name, sampler in worlds.items():
        data = []
        for a, b in itertools.combinations(KEYS, 2):
            for _ in range(per_pair):
                data.append(sampler(a, b))
        g2, df = goodness_of_fit(data, KEYS)
        bad, checked = sst_violations(data, KEYS)
        verdict = "拒绝 BT" if g2 > 32.7 else "不能拒绝 BT"
        print(f"  {pad(name, 34)} {g2:>8.1f} {df:>6d} {pad(verdict, 20, 'right')} {bad:>5d} / {checked:<5d}")
    print("""
观察：
  · 数据确实来自 BT 时，G² 和自由度（21）差不多，检验不报警；SST 也几乎没有明显的违反。
  · 两群口味相反的人混在一起，每群人自己都服从 BT，合起来却不服从：G² 远大于临界值，SST 也有明显违反。
    第 1 讲代码 04 看到的“BT 的预测在几对上都偏了”，在这里被一个正式的检验抓到了。
  · 注意检验需要**足够多**的数据：每对只比较几次时，G² 的卡方近似不可靠，SST 的违反也被噪声淹没。
    真实的 RLHF 数据几乎每一对回答只比较一次，所以这种检验通常要在“回答的特征”层面做（第 5 节），
    或者在大规模的排行榜数据上做。""")


def _mixture_utils(rng, group_a, group_b, share_a=0.6):
    """随机抽一个人（60% 来自 group_a），返回他心里的 8 个效用。"""
    g = group_a if rng.random() < share_a else group_b
    return {k: g.utility(BY_KEY[k]) for k in KEYS}


# =============================================================================
# 演示 3：平局
# =============================================================================
def rk_probs(sa, sb, eta):
    """
    Rao–Kupper（1967）模型：人只有在"差得够多"时才分得出高低，否则说"平局"。

        P(a ≻ b) = σ(s_a − s_b − η)，P(b ≻ a) = σ(s_b − s_a − η)，P(平局) = 1 − 两者之和
    η ≥ 0 是"察觉阈值"：η = 0 时没有平局，就是普通的 BT。（和演示 1 的"咖啡里的糖"是同一个想法。）
    【返回】(P(a≻b), P(b≻a), P(平局))
    """
    pab = sigmoid(sa - sb - eta)
    pba = sigmoid(sb - sa - eta)
    return pab, pba, 1.0 - pab - pba


def davidson_probs(sa, sb, nu):
    """
    Davidson（1970）模型：P(a≻b) ∝ e^{s_a}，P(b≻a) ∝ e^{s_b}，P(平局) ∝ ν·√(e^{s_a} e^{s_b})。

    【思路】平局的"权重"是两者的几何平均乘以 ν：两边实力越接近，平局越常见（相对而言）。
        ν = 0 时退化为 BT。它的好处是保留了"胜负之比 = e^{s_a − s_b}"这个 BT 的性质。
    【返回】(P(a≻b), P(b≻a), P(平局))
    """
    m = max(sa, sb)
    ea, eb = math.exp(sa - m), math.exp(sb - m)
    tie = nu * math.sqrt(ea * eb)
    z = ea + eb + tie
    return ea / z, eb / z, tie / z


def fit_tie_model(table, keys, model):
    """
    用数值梯度上升拟合平局模型。

    【参数】
        table：dict {(a, b): (a 赢次数, b 赢次数, 平局次数)}，a < b
        keys ：选项列表
        model："rk" 或 "davidson"
    【返回】(分数 dict（平均值为 0）, 平局参数 η 或 ν)
    【思路】参数是 M 个分数 + 1 个平局参数（取对数，保证它是正数）。目标是平均对数似然
        （再加一点 L2 让分数的平均值不乱跑）。参数很少，所以用 choice_data.numeric_gradient 就够了。
    """
    n = sum(sum(v) for v in table.values())

    def loglik(params):
        s = dict(zip(keys, params[:-1]))
        extra = math.exp(params[-1])
        total = 0.0
        for (a, b), (wa, wb, t) in table.items():
            pa, pb, pt = (rk_probs if model == "rk" else davidson_probs)(s[a], s[b], extra)
            total += wa * math.log(pa) + wb * math.log(pb) + t * math.log(max(pt, 1e-300))
        return total / n - 1e-4 * math.fsum(x * x for x in params[:-1])

    params = gradient_ascent(loglik, [0.0] * len(keys) + [math.log(0.5)], lr=1.0, iters=4000, tol=1e-7)
    s = dict(zip(keys, params[:-1]))
    m = mean(list(s.values()))
    return {k: v - m for k, v in s.items()}, math.exp(params[-1])


def demo_ties(rng, n=6000, eta=0.8):
    print_header("演示 3：平局 —— 标注员说“差不多”怎么办？")
    keys = ["A", "C", "D", "F", "H"]
    truth = {k: TRUE[k] - mean([TRUE[j] for j in keys]) for k in keys}
    table = {}
    for _ in range(n):
        a, b = sorted(rng.sample(keys, 2))
        pab, pba, _ = rk_probs(TRUE[a], TRUE[b], eta)
        r = rng.random()
        wa, wb, t = table.get((a, b), (0, 0, 0))
        if r < pab:
            wa += 1
        elif r < pab + pba:
            wb += 1
        else:
            t += 1
        table[(a, b)] = (wa, wb, t)
    ties = sum(v[2] for v in table.values())
    print(f"专家按 Rao–Kupper 模型（察觉阈值 η = {eta}）比较 {n} 次，其中 {ties} 次（{ties / n:.0%}）说“差不多”。\n")

    # 方法 1：丢掉平局
    drop = []
    for (a, b), (wa, wb, t) in table.items():
        drop += [pair_to_event(a, b)] * wa + [pair_to_event(b, a)] * wb
    s_drop = fit_choice_model(drop, keys, tol=1e-9)
    # 方法 2：平局算各赢半场（把每次胜负复制两份、每次平局两边各记一次，权重就是 1 : 0.5）
    half = []
    for (a, b), (wa, wb, t) in table.items():
        half += [pair_to_event(a, b)] * (2 * wa + t) + [pair_to_event(b, a)] * (2 * wb + t)
    s_half = fit_choice_model(half, keys, tol=1e-9)
    s_rk, eta_hat = fit_tie_model(table, keys, "rk")
    s_dv, nu_hat = fit_tie_model(table, keys, "davidson")

    print(f"  {'回答':4s} {'真实':>6s} {pad('丢掉平局', 9, 'right')} {pad('各算半场', 9, 'right')} "
          f"{pad('Rao–Kupper', 11, 'right')} {pad('Davidson', 9, 'right')}")
    for k in keys:
        print(f"  {k:4s} {truth[k]:>6.2f} {s_drop[k]:>9.2f} {s_half[k]:>9.2f} {s_rk[k]:>11.2f} {s_dv[k]:>9.2f}")
    def spread(s):
        return max(s.values()) - min(s.values())
    print(f"  {'最高−最低':6s} {spread(truth):>4.2f} {spread(s_drop):>9.2f} {spread(s_half):>9.2f} {spread(s_rk):>11.2f} {spread(s_dv):>9.2f}")
    print(f"\n  Rao–Kupper 估计的察觉阈值 η = {eta_hat:.2f}（真实 {eta}）；Davidson 估计的 ν = {nu_hat:.2f}")
    print("""
观察：
  · 四种做法给出的**顺序**都对，差别在“尺度”上：
    · 丢掉平局：剩下的都是“分得出高低”的比较，看起来比真实的更一边倒，分数差被**放大**了。
    · 各算半场：平局把胜率往 0.5 拉，分数差被**压缩**了。
    · Rao–Kupper 正是生成数据的模型，分数和阈值都估得准。
    · Davidson 对平局的假设和生成数据的方式不同，尺度反而偏大：用错了平局模型，尺度同样会错。
  · 尺度为什么重要？分数差决定了预测的胜率（σ(Δ)），也决定了 RLHF 里奖励模型给出的“奖励差”有多大。
    尺度错了，概率就不校准，下游用 β 调节 KL 约束时也会受影响（第 5 讲：校准；第 8 讲：β 的含义）。
  · 真实数据里平局很常见：Llama 2 让标注员选“明显更好 / 更好 / 稍好 / 差不多（不确定）”；Chatbot Arena 的
    投票里有“平局”和“两个都不好”。常见的简单做法是丢掉平局或算各赢半场，这时要意识到它们对尺度的影响。""")


def main():
    rng = random.Random(1929)        # 1929：Zermelo 发表他的锦标赛排名论文的年份
    demo_graph(rng)
    demo_goodness(rng)
    demo_ties(rng)


if __name__ == "__main__":
    main()