"""
=============================================================================
CS486 · 第 2 讲 · 代码 02：LLM 当裁判——位置偏差、冗长偏差，以及用少量人类标注做无偏校正
=============================================================================

【这个脚本要回答的问题】（讲义 §5.3）
    开放式任务没有标准答案，研究里常用"强 LLM 当裁判"比较两个系统。裁判便宜，但有偏差。
    在一个知道真相的模拟世界里（judge_sim.py）：系统 A 真实地稍好一些，系统 B 更啰嗦；
    裁判偏爱排在前面的回答和更长的回答。

    演示 1：一次评测——单次判、交换位置平均、长度控制、只用人类标注、PPI，各自估出的"A 的胜率"。
            裁判会不会把结论弄反？
    演示 2：重复 300 次评测——每种方法的偏差、波动、95% 区间的覆盖率和宽度。
            "只用裁判"的区间为什么覆盖率是 0？PPI 一定比只用人类标注好吗？
    演示 3：如果长度本身就是质量的一部分（人类也偏爱更完整的回答），长度控制会怎样？
    演示 4：裁判越好，PPI 省下的人类标注越多吗？上限在哪里？

【运行】
    python3 02_llm_judge.py           # 约 10 秒
    python3 02_llm_judge.py --quick   # 冒烟测试（重复次数很少）

【依赖】numpy + 本目录的 judge_sim.py + 第 1 讲的 console.py。
"""

import argparse
import os
import sys
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

import judge_sim as js                                                   # noqa: E402
from console import print_header, print_observations, print_table        # noqa: E402

N_HUMAN = 100     # 人类标注的数量（裁判标注全部 800 个提示）


def demo1_single(seed=7):
    print_header("演示 1：一次评测（800 个提示；人类只标注其中 100 个）——各种方法估出的'A 的胜率'")
    rng = np.random.default_rng(seed)
    world = js.make_world(rng)
    a, b = js.judge_votes(rng, world)
    probs = js.judge_probs(world)
    idx, h = js.human_labels(rng, world, N_HUMAN)
    lc, theta = js.length_controlled(a, b, world["log_ratio"])
    rows = [
        ["真值：人类眼中 A 的胜率 θ", f"{world['theta']:.3f}", ""],
        ["裁判，A 总在前", f"{a.mean():.3f}", "位置偏差偏向 A"],
        ["裁判，B 总在前", f"{b.mean():.3f}", "位置偏差偏向 B"],
        ["裁判，两种顺序平均（交换平均）", f"{((a + b) / 2).mean():.3f}", "去掉了位置偏差，冗长偏差还在"],
        ["长度控制（回归后令长度差 = 0）", f"{lc:.3f}", f"长度系数 θ1 = {theta[1]:+.2f}，位置系数 θ2 = {theta[2]:+.2f}"],
        [f"只用 {N_HUMAN} 个人类标注", f"{h.mean():.3f}", "无偏，但只有 100 个样本"],
        ["PPI（裁判概率分数 + 人类标注）", f"{js.ppi_mean(probs, idx, h)[0]:.3f}", "无偏，利用了全部 800 个裁判分数"],
    ]
    print_table(["方法", "估计的 A 胜率", "说明"], rows, [34, 14, 40], ["<", ">", "<"])
    return {"theta": world["theta"], "a": float(a.mean()), "b": float(b.mean()),
            "swap": float(((a + b) / 2).mean()), "lc": lc, "theta1": float(theta[1]), "theta2": float(theta[2]),
            "ppi": js.ppi_mean(probs, idx, h)[0]}


def simulate(reps, world_kw=None, seed=0):
    """
    重复 reps 次独立的评测，收集每种方法的估计和区间。

    【返回】dict：方法名 → {"err": [估计 − 真值], "cover": 覆盖次数, "width": [区间宽度]}
        没有区间的方法（单次判、长度控制）cover 和 width 为空。
    """
    rng = np.random.default_rng(seed)
    out = {}

    def add(name, est, theta, lo=None, hi=None):
        d = out.setdefault(name, {"err": [], "cover": 0, "width": []})
        d["err"].append(est - theta)
        if lo is not None:
            d["cover"] += lo <= theta <= hi
            d["width"].append(hi - lo)

    for _ in range(reps):
        world = js.make_world(rng, world_kw)
        th = world["theta"]
        a, b = js.judge_votes(rng, world)
        probs = js.judge_probs(world)
        idx, h = js.human_labels(rng, world, N_HUMAN)
        vote = (a + b) / 2
        add("裁判，A 总在前", float(a.mean()), th)
        m, lo, hi = js.judge_only_ci(vote)
        add("裁判，交换平均（只用裁判的区间）", m, th, lo, hi)
        add("长度控制", js.length_controlled(a, b, world["log_ratio"])[0], th)
        m, lo, hi = js.human_only(h)
        add(f"只用 {N_HUMAN} 个人类标注", m, th, lo, hi)
        e, lo, hi, _ = js.ppi_mean(vote, idx, h, lam=1.0)
        add("PPI，裁判投票（λ = 1）", e, th, lo, hi)
        e, lo, hi, _ = js.ppi_mean(vote, idx, h, lam=None)
        add("PPI++，裁判投票（估计 λ）", e, th, lo, hi)
        e, lo, hi, _ = js.ppi_mean(probs, idx, h, lam=1.0)
        add("PPI，裁判概率分数（λ = 1）", e, th, lo, hi)
    return out


def summarize(out, reps):
    rows = []
    for name, d in out.items():
        err = np.array(d["err"])
        cov = f"{100 * d['cover'] / reps:.1f}%" if d["width"] else "-"
        wid = f"{np.mean(d['width']):.3f}" if d["width"] else "-"
        rows.append([name, f"{err.mean():+.3f}", f"{err.std():.3f}", cov, wid])
    print_table(["方法", "平均偏差", "波动（标准差）", "95% 区间覆盖率", "区间平均宽度"], rows,
                [34, 10, 14, 16, 14], ["<", ">", ">", ">", ">"])


def demo2_repeated(reps):
    print_header(f"演示 2：重复 {reps} 次独立评测——偏差、波动与置信区间")
    out = simulate(reps)
    summarize(out, reps)
    return out


def demo3_length_is_quality(reps):
    print_header(f"演示 3：如果长度本身就是质量的一部分（人类也偏爱更长的回答：human_len = 0.8）——重复 {reps} 次")
    out = simulate(reps, {"human_len": 0.8}, seed=1)
    keep = {k: out[k] for k in ("裁判，A 总在前", "裁判，交换平均（只用裁判的区间）", "长度控制",
                                f"只用 {N_HUMAN} 个人类标注", "PPI，裁判概率分数（λ = 1）")}
    summarize(keep, reps)
    return keep


def demo4_judge_quality(reps):
    print_header(f"演示 4：裁判越好，PPI 的区间越窄吗？（裁判概率分数；每行重复 {reps} 次）")
    rows, out = [], {}
    settings = [("原来的裁判（judge_beta = 0.8）", {"judge_beta": 0.8}),
                ("更敏锐的裁判（judge_beta = 2）", {"judge_beta": 2.0}),
                ("更敏锐的裁判（judge_beta = 4）", {"judge_beta": 4.0}),
                ("无偏且敏锐（judge_beta = 1，无长度 / 位置偏差）", {"judge_beta": 1.0, "judge_len": 0.0, "judge_pos": 0.0})]
    for label, kw in settings:
        rng = np.random.default_rng(3)
        w_ppi, w_pp, w_h, agree_jh, agree_hh = [], [], [], [], []
        for _ in range(reps):
            world = js.make_world(rng, kw)
            probs = js.judge_probs(world)
            idx, h = js.human_labels(rng, world, N_HUMAN)
            _, lo, hi, _ = js.ppi_mean(probs, idx, h)
            w_ppi.append(hi - lo)
            _, lo, hi, _ = js.ppi_mean(probs, idx, h, lam=None)
            w_pp.append(hi - lo)
            _, lo2, hi2 = js.human_only(h)
            w_h.append(hi2 - lo2)
            a, _ = js.judge_votes(rng, world)
            h1 = (rng.random(len(probs)) < world["p_human"]).astype(float)
            h2 = (rng.random(len(probs)) < world["p_human"]).astype(float)
            agree_jh.append(float((a == h1).mean()))
            agree_hh.append(float((h1 == h2).mean()))
        ratio = float(np.mean(w_ppi) / np.mean(w_h))
        ratio_pp = float(np.mean(w_pp) / np.mean(w_h))
        out[label] = (ratio, float(np.mean(agree_jh)), float(np.mean(agree_hh)), ratio_pp)
        rows.append([label, f"{np.mean(agree_jh):.3f}", f"{np.mean(agree_hh):.3f}",
                     f"{ratio:.3f}", f"{ratio_pp:.3f}", f"{1 / ratio ** 2:.2f}×"])
    # 上限：一个"读心"的裁判，直接给出 p_human
    rng = np.random.default_rng(4)
    w_ppi, w_h = [], []
    for _ in range(reps):
        world = js.make_world(rng)
        idx, h = js.human_labels(rng, world, N_HUMAN)
        _, lo, hi, _ = js.ppi_mean(world["p_human"], idx, h)
        w_ppi.append(hi - lo)
        _, lo2, hi2 = js.human_only(h)
        w_h.append(hi2 - lo2)
    ratio = float(np.mean(w_ppi) / np.mean(w_h))
    out["上限"] = (ratio, None, None, None)
    rows.append(["上限：裁判分数 = 人类偏好的真实概率", "-", "-", f"{ratio:.3f}", "-", f"{1 / ratio ** 2:.2f}×"])
    print_table(["裁判", "裁判–人类一致率", "人类–人类一致率", "宽度比 PPI", "宽度比 PPI++", "相当于几倍人类标注"],
                rows, [44, 16, 16, 12, 14, 18], ["<", ">", ">", ">", ">", ">"])
    print("（宽度比 = PPI 区间宽 / 只用人类标注的区间宽；'相当于几倍人类标注'= (λ = 1 的宽度比)^(-2)，"
          "因为区间宽度和 1/√n 成正比。）")
    return out


def main():
    ap = argparse.ArgumentParser(description="LLM 当裁判的偏差与校正")
    ap.add_argument("--quick", action="store_true", help="冒烟测试：重复次数很少")
    args = ap.parse_args()
    t0 = time.perf_counter()
    reps = 20 if args.quick else 300

    one = demo1_single()
    rep = demo2_repeated(reps)
    lq = demo3_length_is_quality(reps)
    jq = demo4_judge_quality(40 if args.quick else 200)

    if args.quick:
        print("\n（--quick 冒烟测试：重复次数太少，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return

    def mean_err(d, k):
        return float(np.mean(d[k]["err"]))

    def width(d, k):
        return float(np.mean(d[k]["width"]))

    hum = f"只用 {N_HUMAN} 个人类标注"
    K0, K2, K4 = "原来的裁判（judge_beta = 0.8）", "更敏锐的裁判（judge_beta = 2）", "更敏锐的裁判（judge_beta = 4）"
    print_observations([
        f"演示 1：真实胜率 {one['theta']:.3f}（A 更好），但裁判在 A 在前时说 {one['a']:.3f}、B 在前时说 {one['b']:.3f}；"
        f"交换位置平均后是 {one['swap']:.3f}——**结论翻转了**：裁判认为 B 更好。原因是 B 更长，裁判有冗长偏差。"
        f"长度控制回归把长度的影响拿掉后是 {one['lc']:.3f}，回到了正确的一侧。"
        f"（这一次 PPI 的估计是 {one['ppi']:.3f}：PPI 是无偏的，但单次评测仍有约 ±0.08 的波动——见演示 2 的标准差。）",
        f"演示 2：'只用裁判'的区间覆盖率是 {100 * rep['裁判，交换平均（只用裁判的区间）']['cover'] / reps:.1f}%——"
        f"它的平均偏差 {mean_err(rep, '裁判，交换平均（只用裁判的区间）'):+.3f} 远大于区间宽度。"
        "置信区间只反映'抽样噪声'，不反映'测量本身的偏差'；评测的样本再多，也治不了有偏的裁判。",
        f"演示 2：只用 {N_HUMAN} 个人类标注是无偏的（覆盖率 {100 * rep[hum]['cover'] / reps:.1f}%），但区间宽 {width(rep, hum):.3f}。"
        f"经典 PPI 用裁判的 0/1 投票时区间反而**更宽**（{width(rep, 'PPI，裁判投票（λ = 1）'):.3f}）："
        "裁判投票和人类判断的一致性不够高，校正项的方差比省下的方差还大。"
        f"PPI++ 估计 λ 之后是 {width(rep, 'PPI++，裁判投票（估计 λ）'):.3f}，"
        f"改用裁判的概率分数是 {width(rep, 'PPI，裁判概率分数（λ = 1）'):.3f}——三种 PPI 的覆盖率都在 95% 附近，"
        "都是无偏的；但能省下多少人类标注，取决于裁判分数和人类判断有多相关。",
        f"演示 3：如果人类本身也偏爱更长、更完整的回答，长度控制就会**过度校正**（平均偏差 "
        f"{mean_err(lq, '长度控制'):+.3f}），而 PPI 仍然无偏（{mean_err(lq, 'PPI，裁判概率分数（λ = 1）'):+.3f}）。"
        "长度控制依赖一个假设——'长度不是质量的一部分'；PPI 只依赖'人类标注是随机抽的'。",
        f"演示 4：更'敏锐'（但仍有长度和位置偏差）的裁判和人类的一致率更高（{jq[K0][1]:.3f} → {jq[K4][1]:.3f}），"
        f"λ = 1 的 PPI 却越来越差（宽度比 {jq[K0][0]:.3f} → {jq[K2][0]:.3f} → {jq[K4][0]:.3f}）：分数过于自信又带着偏差时，"
        f"校正项 h − ĵ 的方差变大。PPI++ 把 λ 缩小后是 {jq[K0][3]:.3f} / {jq[K2][3]:.3f} / {jq[K4][3]:.3f}，"
        "不再比只用人类标注差。决定 PPI 收益的是裁判分数和'人类偏好的概率'有多接近（校准），而不是一致率。",
        f"即使裁判分数恰好等于人类偏好的真实概率，区间也只缩到人类区间的 {jq['上限'][0]:.3f}"
        f"（相当于 {1 / jq['上限'][0] ** 2:.2f} 倍人类标注）。上限来自人类标注本身的随机性："
        f"这个世界里两次人类判断只有约 {100 * jq[K0][2]:.0f}% 一致——单个人类标注本来就是一个噪声很大的测量。"
        "真实场景中人类分歧小、裁判又校准得好时，PPI 能省下的标注会多得多。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()