"""
=============================================================================
第 14 讲 · 代码 1：风格控制 —— 把"内容"和"风格"分开（以及什么时候会分错）
=============================================================================

对应讲义：lectures/lecture14_style_surrogates.md 第 2、3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
模拟一个有 12 个模型的竞技场。每个模型有真实的内容质量，也有自己的"风格习惯"：有的回答又长又爱用标题和列表，
有的简短朴素。投票者除了看内容，还偏爱长回答、偏爱格式（Chatbot Arena 的分析发现真实用户也是这样）。

    演示 1：普通 BT vs 加了风格协变量的 BT：谁更接近真实的内容质量？风格的影响有多大？
    演示 2："坏控制"：如果长回答**真的**更好（内容更完整），控制长度就把真实的好处也减掉了。
            回归分不清"长度是偏差"还是"长度是质量" —— 这要靠人来判断。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_style_control.py        # 约 3 秒

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

import numpy as np

from style_tools import bt_design, fit_logistic, kendall_tau, pad, print_header, simulate_arena


def scores_and_style(arena, control):
    """拟合 BT（control=True 时加风格协变量），返回 (减去平均值的模型分数, 风格系数, 风格系数的标准误)。"""
    X = bt_design(arena["i"], arena["j"], arena["M"], arena["style_diff"] if control else None)
    beta, cov = fit_logistic(X, arena["y"])
    M = arena["M"]
    s = beta[:M] - beta[:M].mean()
    if not control:
        return s, None, None
    return s, beta[M:], np.sqrt(np.diag(cov)[M:])


def ranks(x):
    """名次（1 = 最高）。"""
    order = np.argsort(-x, kind="stable")
    r = np.empty(len(x), dtype=int)
    r[order] = np.arange(1, len(x) + 1)
    return r


def demo_style_control(rng):
    print_header("演示 1：普通 BT vs 风格控制")
    arena = simulate_arena(rng, n_battles=20000)
    q = arena["quality"] - arena["quality"].mean()
    s_naive, _, _ = scores_and_style(arena, control=False)
    s_ctrl, gamma, se = scores_and_style(arena, control=True)
    print(f"12 个模型、{len(arena['y'])} 场对战。投票者的风格偏爱：长度 γ = 1.0、格式 γ = 0.6（乘的是归一化差）。\n")
    print(f"  {pad('模型', 5)} {pad('真实质量', 9, 'right')} {pad('平均长度', 9, 'right')} {pad('格式', 5, 'right')} "
          f"{pad('普通 BT', 8, 'right')} {pad('名次', 5, 'right')} {pad('风格控制', 9, 'right')} {pad('名次', 5, 'right')} "
          f"{pad('真实名次', 9, 'right')}")
    rq, rn, rc = ranks(q), ranks(s_naive), ranks(s_ctrl)
    for m in range(arena["M"]):
        print(f"  {pad(f'M{m}', 5)} {q[m]:>+9.2f} {arena['mean_len'][m]:>9.0f} {arena['md_rate'][m]:>5.0f} "
              f"{s_naive[m]:>+8.2f} {rn[m]:>5d} {s_ctrl[m]:>+9.2f} {rc[m]:>5d} {rq[m]:>9d}")
    print(f"\n  和真实质量的肯德尔 τ：普通 BT {kendall_tau(s_naive, q):.3f}，风格控制 {kendall_tau(s_ctrl, q):.3f}")
    print(f"  估计的风格系数：长度 {gamma[0]:.3f} ± {1.96 * se[0]:.3f}，格式 {gamma[1]:.3f} ± {1.96 * se[1]:.3f}"
          f"（真实值 1.0、0.6）")
    print(f"  普通 BT 的分数和模型平均长度的相关系数：{np.corrcoef(s_naive, np.log(arena['mean_len']))[0, 1]:.3f}；"
          f"风格控制之后：{np.corrcoef(s_ctrl, np.log(arena['mean_len']))[0, 1]:.3f}")
    print("""
观察：
  · 普通 BT 把“写得长、格式多”算进了实力：又长又爱用格式的模型（M1、M3、M5、M7、M9）被抬高，
    简短朴素的模型（M0、M2、M6、M8）被压低。普通 BT 的分数和平均长度明显正相关。
  · 风格控制把长度和格式的差作为协变量放进 BT，模型分数就只剩“控制了风格之后”的部分 ——
    排序和真实的内容质量几乎一致，风格系数也估计得很准。
  · 这就是 Chatbot Arena 2024 年推出的“风格控制”排行榜：同一批投票，控制风格之后，一些模型的名次变化很大。
  · 原理和第 2 讲的“遗漏变量偏差”一样：风格既影响投票、又和模型有关（某些模型总写得长），
    不放进模型，它的影响就被算到模型头上。""")


def demo_bad_control(rng):
    print_header("演示 2：坏控制 —— 如果长回答真的更好")
    arena = simulate_arena(rng, n_battles=20000, gamma_len=0.0, gamma_md=0.0, length_is_quality=1.0)
    # "真实的综合质量"：每个模型对随机对手的期望胜率（用一大批新对战模拟，内容质量里包含了长度的真实贡献）
    big = simulate_arena(np.random.default_rng(99), n_battles=200000, gamma_len=0.0, gamma_md=0.0, length_is_quality=1.0)
    win = np.array([np.mean(np.where(big["i"] == m, big["y"], 1 - big["y"])[(big["i"] == m) | (big["j"] == m)])
                    for m in range(big["M"])])
    s_naive, _, _ = scores_and_style(arena, control=False)
    s_ctrl, gamma, se = scores_and_style(arena, control=True)
    print("这一次投票者**没有**风格偏爱；可是长回答的内容确实更完整（长度的归一化差对内容质量的贡献系数 = 1.0）。")
    print("“真实的综合质量”用每个模型对随机对手的期望胜率衡量。\n")
    print(f"  {pad('', 28)} {pad('和综合质量的肯德尔 τ', 21, 'right')}")
    print(f"  {pad('普通 BT', 28)} {kendall_tau(s_naive, win):>21.3f}")
    print(f"  {pad('控制了长度和格式的 BT', 28)} {kendall_tau(s_ctrl, win):>21.3f}")
    print(f"\n  估计的“长度”系数：{gamma[0]:.3f} ± {1.96 * se[0]:.3f} —— 回归发现了长度的影响，")
    print("  但它不知道这个影响是“投票者的偏差”（演示 1）还是“长回答真的更好”（这里）。")
    print("""
观察：
  · 这里普通 BT 才是对的：长度带来的胜率是真实的内容价值。控制长度之后，长回答的模型被“扣分”，
    排序反而偏离了真实的综合质量。
  · 统计上，演示 1 和演示 2 的数据**看起来一样**：长回答更容易赢。回归只能告诉你“长度和胜负有关”，
    不能告诉你这种关系是“偏差”还是“价值”。这是一个因果问题，要靠领域知识、或者专门设计的实验
    （例如把同一个回答改写成长短两个版本、内容不变，看投票者怎么选）来回答。
  · 风格控制的价值判断：Chatbot Arena 同时公布控制前和控制后的排行榜，把选择留给读者。
    第 2 讲的“遗漏变量偏差”和这里的“坏控制”是一对：漏掉该控制的变量有偏差，控制了不该控制的变量也有偏差。""")


def main():
    rng = np.random.default_rng(14)
    demo_style_control(rng)
    demo_bad_control(rng)


if __name__ == "__main__":
    main()