"""
=============================================================================
第 17 讲 · 01：我的方法真的更好吗？—— 配对检验、多次运行、"在测试集上挑最好的"
=============================================================================

对应讲义：lectures/lecture17_future_research.md 第 3 节

-----------------------------------------------------------------------------
【设定】
-----------------------------------------------------------------------------
一个 benchmark 有 N 道题，题目的难度差别很大（难度 θ ~ N(0, 3²)：有很多题谁都做不出、也有很多题谁都能做出）。
agent 在一道题上成功的概率 = sigmoid(能力 − θ)。agent A 的能力 0.5；agent B 比 A 强一点（能力 0.75），
平均成功率大约高 3 个百分点。每个 agent 在每道题上跑一次，结果是成功或失败。

演示 1：同样的数据，两种检验 —— 不配对（把两个 agent 的成功率当作两组独立的样本比较）vs
        配对（McNemar 检验：只看"A 对 B 错"和"A 错 B 对"的题）。重复 400 次评估，看各自多常检测出差别。
演示 2：同一个 agent、同一套题，跑 5 次，分数差多少？
演示 3：20 个"改进"的提示词，其实和基线一模一样。在测试集上挑表现最好的那个报告，会怎样？

所有设定都是模拟的，演示的是统计上的道理。

运行时间：约 2 秒（只用标准库）。
"""

import math
import random

THETA_SD = 3.0
ABILITY_A, ABILITY_B = 0.5, 0.75
N_EVALS = 400               # 演示 1：重复多少次"评估"来估计检出率


def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z))


def outcomes(ability, thetas, rng):
    """agent 在每道题上跑一次：成功 1、失败 0。"""
    return [1 if rng.random() < sigmoid(ability - t) else 0 for t in thetas]


def unpaired_z(a, b):
    """
    不配对的两比例 z 检验。

    【思路】把 a、b 当作两组独立的伯努利样本：z = (p̂_B − p̂_A) / √(p̄(1 − p̄)(2/N))，p̄ 是合并的成功率。
    【边界】p̄ 是 0 或 1 时分母为 0，返回 0（没有差别可言）。
    """
    n = len(a)
    pa, pb = sum(a) / n, sum(b) / n
    pbar = (pa + pb) / 2
    se = math.sqrt(pbar * (1 - pbar) * 2 / n)
    return 0.0 if se == 0 else (pb - pa) / se


def paired_z(a, b):
    """
    配对检验（McNemar）：只看两个 agent 结果不同的题。

    【思路】同一道题上，两个 agent 都对或都错，不提供"谁更好"的信息。设 b10 = A 错 B 对 的题数，
        b01 = A 对 B 错 的题数。如果两者一样好，每道"结果不同"的题偏向 B 的概率是 1/2，
        z = (b10 − b01) / √(b10 + b01)。
    【为什么更有力】题目难度的差别是两个 agent 共有的：难题两个都错，易题两个都对。不配对的检验把这些
        差别当成"噪声"，配对检验把它们抵消掉了。
    """
    b10 = sum(1 for x, y in zip(a, b) if x == 0 and y == 1)
    b01 = sum(1 for x, y in zip(a, b) if x == 1 and y == 0)
    return 0.0 if b10 + b01 == 0 else (b10 - b01) / math.sqrt(b10 + b01)


def demo_paired():
    print("=" * 78)
    print("演示 1：检出率（|z| > 1.96，即 p < 0.05）—— 不配对 vs 配对")
    print("=" * 78)
    rng = random.Random(1700)
    big = [rng.gauss(0, THETA_SD) for _ in range(200000)]
    gap = sum(sigmoid(ABILITY_B - t) - sigmoid(ABILITY_A - t) for t in big) / len(big)
    print(f"B 比 A 的真实平均成功率高 {gap * 100:.1f} 个百分点")
    print(f"{'题数 N':>8}{'不配对的检出率':>16}{'配对的检出率':>14}{'A = B 时配对的误报率':>24}")
    for n in (100, 300, 1000, 3000):
        hit_u = hit_p = false_p = 0
        for _ in range(N_EVALS):
            thetas = [rng.gauss(0, THETA_SD) for _ in range(n)]
            a, b = outcomes(ABILITY_A, thetas, rng), outcomes(ABILITY_B, thetas, rng)
            hit_u += abs(unpaired_z(a, b)) > 1.96
            hit_p += abs(paired_z(a, b)) > 1.96
            a2 = outcomes(ABILITY_A, thetas, rng)                    # 和 A 一模一样的 agent
            false_p += abs(paired_z(a, a2)) > 1.96
        print(f"{n:8d}{hit_u / N_EVALS:16.3f}{hit_p / N_EVALS:14.3f}{false_p / N_EVALS:24.3f}")
    print("\n观察：")
    print("  · 差 3 个百分点，在 1000 道题上用不配对的检验只有 14% 的机会检测出来，配对检验有 44%；")
    print("    3000 道题时 67% vs 92%。同样的数据，换一种检验，检出率提高好几倍。")
    print("  · 原因：题目难度的差别是两个 agent 共有的（难题都错、易题都对），配对检验把它抵消了。")
    print("    比较两个 agent 时，一定要在**同一批题**上跑，并且用配对的方法（McNemar、配对 bootstrap）。")
    print("  · 两个 agent 一模一样时，配对检验的误报率在 5% 附近 —— 它没有因为更灵敏而变得\"乱报\"。")
    print("  · 反过来说：几百道题的 benchmark 上，差 2~3 个点的结论，即使检验方法正确，也常常检测不出来。\n")


def demo_runs():
    print("=" * 78)
    print("演示 2：同一个 agent、同一套 200 道题，跑 5 次")
    print("=" * 78)
    rng = random.Random(1710)
    thetas = [rng.gauss(0, THETA_SD) for _ in range(200)]
    scores = [sum(outcomes(ABILITY_A, thetas, rng)) / len(thetas) for _ in range(5)]
    mean = math.fsum(scores) / len(scores)
    sd = math.sqrt(math.fsum((s - mean) ** 2 for s in scores) / (len(scores) - 1))
    print("  5 次的成功率：" + "、".join(f"{s:.3f}" for s in scores))
    print(f"  平均 {mean:.3f}，标准差 {sd:.3f}，最高和最低相差 {max(scores) - min(scores):.3f}")
    print("\n观察：")
    print("  · 什么都没改，同一套题跑 5 次，分数就相差 3 个百分点：agent 的输出有随机性（采样温度、工具的返回、")
    print("    环境的变化）。只跑一次就说\"我的方法高了 2 个点\"，可能只是运气。")
    print("  · 报告多次运行的平均和标准差（或置信区间）；比较两个方法时，每个方法都跑多次。\n")


def demo_forking_paths():
    print("=" * 78)
    print("演示 3：20 个和基线一模一样的\"改进\"，在 200 道题的测试集上挑最好的报告（重复 400 次）")
    print("=" * 78)
    rng = random.Random(1720)
    claimed, significant, honest = [], 0, []
    for _ in range(N_EVALS):
        test = [rng.gauss(0, THETA_SD) for _ in range(200)]
        val = [rng.gauss(0, THETA_SD) for _ in range(200)]
        base = outcomes(ABILITY_A, test, rng)
        variants_test = [outcomes(ABILITY_A, test, rng) for _ in range(20)]
        best = max(range(20), key=lambda i: sum(variants_test[i]))
        claimed.append((sum(variants_test[best]) - sum(base)) / len(test))
        significant += paired_z(base, variants_test[best]) > 1.96
        variants_val = [sum(outcomes(ABILITY_A, val, rng)) for _ in range(20)]
        chosen = max(range(20), key=lambda i: variants_val[i])         # 在验证集上挑
        honest.append((sum(variants_test[chosen]) - sum(base)) / len(test))
    print(f"  在测试集上挑最好的：平均\"提升\" {math.fsum(claimed) / len(claimed) * 100:+.1f} 个百分点，"
          f"{significant / N_EVALS:.0%} 的情况下\"显著\"（p < 0.05）")
    print(f"  在验证集上挑、在测试集上报告：平均提升 {math.fsum(honest) / len(honest) * 100:+.1f} 个百分点")
    print("\n观察：")
    print("  · 20 个\"改进\"其实都和基线一样好，但在测试集上挑最好的那个，平均\"提升\"超过 4 个百分点，")
    print("    将近四分之一的情况下还是\"显著\"的。这就是\"分叉小径的花园\"（garden of forking paths）：")
    print("    试的东西越多，越容易在噪声里找到一个看起来好的。")
    print("  · 在验证集上挑、在没碰过的测试集上报告，\"提升\"就回到了 0 附近 —— 这才是真实的效果。")
    print("    调提示词、调超参数、挑 checkpoint，都应该在验证集上做；测试集只在最后看一次。\n")


if __name__ == "__main__":
    demo_paired()
    demo_runs()
    demo_forking_paths()