"""
=============================================================================
第 14 讲 · 代码 2：长度控制的胜率 —— AlpacaEval 的做法，以及"把回答写长"能不能刷分
=============================================================================

对应讲义：lectures/lecture14_style_surrogates.md 第 3.3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
AlpacaEval：一组固定的指令，每个被测模型的回答和一个**参考模型**的回答比，由一个 LLM 评委判谁更好，
报告"胜率"。LLM 评委偏爱长回答（第 1、8 讲），于是"把回答写长"就能刷高胜率。

长度控制的胜率（Dubois 等 2024）：对每个模型拟合一个逻辑回归

        logit P(模型 胜 参考) = θ + φ · tanh(长度差 / 标准差)

再报告"长度差为 0 时"的预测胜率 σ(θ) —— "如果两个回答一样长，评委会怎么判"。（原文还加了指令难度项，这里省略。）

    演示 1：10 个模型的原始胜率、长度控制胜率，和"评委没有长度偏好时"的真实胜率比较。
    演示 2：刷分：让一个模型"写得更啰嗦"（内容不变、长度翻倍），两种胜率各变了多少？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_length_controlled_winrate.py        # 约 1 秒

只需要 numpy。依赖本目录的 style_tools.py。
"""

import numpy as np

from style_tools import fit_logistic, kendall_tau, pad, print_header, sigmoid

N_INSTR = 800            # 指令数
REF_LEN = 1000.0         # 参考模型回答的平均长度（字）
LEN_SCALE = 500.0        # tanh 的尺度
JUDGE_BIAS = 0.8         # 评委对长度的偏爱


def evaluate_model(theta, mean_len, rng, bias=JUDGE_BIAS):
    """
    让评委把一个模型在全部指令上的回答和参考回答比较。

    【参数】theta：模型相对参考模型的内容质量（logit）；mean_len：模型回答的平均长度
    【返回】(胜负数组, 长度差数组)
    """
    len_m = mean_len * np.exp(rng.normal(0, 0.3, N_INSTR))
    len_r = REF_LEN * np.exp(rng.normal(0, 0.3, N_INSTR))
    d = len_m - len_r
    win = (rng.random(N_INSTR) < sigmoid(theta + bias * np.tanh(d / LEN_SCALE))).astype(int)
    return win, d


def lc_winrate(win, d):
    """
    长度控制的胜率：拟合 logit = θ + φ·tanh(d/尺度)，返回 σ(θ)（长度差为 0 时的预测胜率）和 φ。
    """
    X = np.stack([np.ones(len(d)), np.tanh(d / LEN_SCALE)], axis=1)
    beta, _ = fit_logistic(X, win)
    return float(sigmoid(beta[0])), float(beta[1])


def demo_lc(rng):
    print_header("演示 1：原始胜率 vs 长度控制的胜率")
    thetas = np.array([0.9, 0.6, 0.5, 0.3, 0.2, 0.0, -0.1, -0.3, -0.5, -0.8])
    lens = np.array([600, 1800, 900, 2200, 1000, 2500, 700, 1600, 800, 2000], dtype=float)
    print(f"{N_INSTR} 条指令；参考模型的回答平均 {REF_LEN:.0f} 字；评委的长度偏爱 φ = {JUDGE_BIAS}（乘 tanh(长度差/{LEN_SCALE:.0f})）。")
    print("“真实胜率”= 评委没有长度偏爱时的期望胜率 σ(θ)。\n")
    print(f"  {pad('模型', 5)} {pad('平均长度', 9, 'right')} {pad('真实胜率', 9, 'right')} {pad('原始胜率', 9, 'right')} "
          f"{pad('长度控制胜率', 13, 'right')} {pad('估计的 φ', 9, 'right')}")
    true_wr, raw, lc = [], [], []
    for m, (th, L) in enumerate(zip(thetas, lens)):
        win, d = evaluate_model(th, L, rng)
        w_lc, phi = lc_winrate(win, d)
        true_wr.append(float(sigmoid(th)))
        raw.append(win.mean())
        lc.append(w_lc)
        print(f"  {pad(f'M{m}', 5)} {L:>9.0f} {true_wr[-1]:>9.3f} {raw[-1]:>9.3f} {w_lc:>13.3f} {phi:>9.2f}")
    print(f"\n  和真实胜率的肯德尔 τ：原始胜率 {kendall_tau(np.array(raw), np.array(true_wr)):.3f}，"
          f"长度控制胜率 {kendall_tau(np.array(lc), np.array(true_wr)):.3f}")
    print(f"  平均绝对误差：原始 {np.mean(np.abs(np.array(raw) - true_wr)):.3f}，"
          f"长度控制 {np.mean(np.abs(np.array(lc) - true_wr)):.3f}")
    print("""
观察：
  · 原始胜率被长度拉偏：长回答的模型（M1、M3、M5、M9）胜率虚高，短回答的模型（M0、M6）被低估。
  · 长度控制的胜率回答的是“如果两个回答一样长，评委会怎么判”，更接近真实胜率，排序也更准。
    大多数模型的 φ 估计在 0.8 附近（模拟里评委对所有模型的长度偏爱相同）。
  · 回答特别长的模型（M3、M5）估计得不准：它们的长度差大多落在 tanh 的饱和区，“长度差为 0”的情况几乎没出现过，
    长度控制胜率靠的是外推，误差明显变大（M5 被低估了约 0.08）。
  · 原文（Dubois 等 2024）报告：长度控制让 AlpacaEval 和 Chatbot Arena（人类投票）的排名相关性明显提高。""")
    return thetas, lens


def demo_gaming(rng, thetas, lens):
    print_header("演示 2：刷分 —— 内容不变，只是写得更啰嗦")
    m = 4
    print(f"模型 M{m}（真实胜率 {sigmoid(thetas[m]):.3f}）。在系统提示里加一句“回答得详细一些”：内容质量不变，长度变成原来的 k 倍。\n")
    print(f"  {pad('长度倍数 k', 11)} {pad('平均长度', 9, 'right')} {pad('原始胜率', 9, 'right')} {pad('长度控制胜率', 13, 'right')}")
    for k in [1.0, 1.5, 2.0, 3.0]:
        res = []
        for rep in range(20):
            win, d = evaluate_model(thetas[m], lens[m] * k, np.random.default_rng(1000 + rep))
            res.append((win.mean(), lc_winrate(win, d)[0]))
        r = np.array(res).mean(axis=0)
        print(f"  {k:>11.1f} {lens[m] * k:>9.0f} {r[0]:>9.3f} {r[1]:>13.3f}")
    print("""
（每一行是 20 次评测的平均。）

观察：
  · 只改长度，原始胜率就能涨近 20 个百分点（0.55 → 0.73）—— 不需要任何真正的改进。
  · 长度控制胜率几乎不动：它把“长度差”的影响单独拿出来，再把它设成 0。
  · 但它不是万能的：k 很大时，几乎所有回答都比参考长，tanh 饱和，“长度差为 0”成了外推（数据里几乎没有这种情况），
    估计的方差变大。原文也指出，长度控制减少了“刷长度”的空间，但模型可以去利用评委的其他偏好（格式、语气、自信）。
  · 任何一个公开的、被大家优化的指标，最终都会被“刷”—— 这就是下一节的古德哈特定律。""")


def main():
    rng = np.random.default_rng(141)
    thetas, lens = demo_lc(rng)
    demo_gaming(rng, thetas, lens)


if __name__ == "__main__":
    main()