"""
=============================================================================
第 17 讲 · 02：公平的基线与消融 —— 一个"新方法"到底好在哪里
=============================================================================

对应讲义：lectures/lecture17_future_research.md 第 3.3、3.4 节

-----------------------------------------------------------------------------
【一篇虚构的论文】
-----------------------------------------------------------------------------
"我们提出 Revise-and-Pick：模型先回答一次，再自己修改两次（第 2 讲的 revise），最后让一个 verifier
从这三个版本里挑一个。和直接回答相比，准确率提高了 X 个百分点！"

这个比较公平吗？Revise-and-Pick 用了 3 次生成 + verifier，直接回答只用了 1 次生成。
本文件做三件事：
    1. 复现论文的比较（vs 1 次生成）；
    2. 加上**计算量相同**的基线：3 个样本多数投票、3 个样本 + 同一个 verifier 挑最好的（best-of-3）、
       weighted best-of-3（第 2 讲）；
    3. **消融**：去掉 verifier（只取最后一次修改的结果）、去掉修改（就是 best-of-3）。
每个比较都报告和 Revise-and-Pick 的配对差值，以及配对 bootstrap 的 95% 区间。

题目和模型来自第 1、2 讲的 toy world：10B 模型、temperature 0.7、500 道 3~6 步的题；verifier 用第 2 讲的
verifier_score（signal = 1.0，一个比较弱的 verifier）。

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

import math
import os
import random
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture02"))

from tts_tools import (ToyLLM, majority_vote, make_problem_set, revise, select_by_verifier,  # noqa: E402
                       select_weighted)

TEMPERATURE = 0.7
SIGNAL = 1.0
N_BOOT = 2000


def paired_bootstrap(xs, ys, rng):
    """
    配对 bootstrap：对题目重抽样，估计 mean(xs) − mean(ys) 的 95% 区间。

    【参数】xs、ys：同一批题上两个方法的对错（0/1），等长
    【思路】每次有放回地抽 N 道题（两个方法用**同一组**题），算差值；重复 N_BOOT 次，取 2.5% 和 97.5% 分位数。
    """
    n = len(xs)
    diffs = []
    for _ in range(N_BOOT):
        idx = [rng.randrange(n) for _ in range(n)]
        diffs.append(sum(xs[i] - ys[i] for i in idx) / n)
    diffs.sort()
    return diffs[int(0.025 * N_BOOT)], diffs[int(0.975 * N_BOOT) - 1]


def main():
    model = ToyLLM(1e10)
    problems = make_problem_set(random.Random(1750), 500, 3, 6)
    rng = random.Random(1751)
    results = {name: [] for name in ("revise_pick", "single", "majority3", "best3", "weighted3", "revise_only")}
    for p in problems:
        chain = [model.generate(p, rng, TEMPERATURE)]
        for _ in range(2):
            chain.append(revise(model, p, chain[-1], rng, TEMPERATURE))
        results["revise_pick"].append(select_by_verifier(p, chain, SIGNAL).correct)
        results["revise_only"].append(chain[-1].correct)
        results["single"].append(chain[0].correct)                       # 论文的基线：第一次的回答
        three = [model.generate(p, rng, TEMPERATURE) for _ in range(3)]
        results["majority3"].append(majority_vote([s.answer for s in three])[0] == p.answer())
        results["best3"].append(select_by_verifier(p, three, SIGNAL).correct)
        results["weighted3"].append(select_weighted(p, three, SIGNAL).correct)

    labels = (("single", "直接回答（1 次生成）—— 论文的基线", 1),
              ("majority3", "3 个样本多数投票", 3),
              ("best3", "3 个样本 + verifier（best-of-3）", 3),
              ("weighted3", "weighted best-of-3", 3),
              ("revise_only", "消融：只修改，不用 verifier", 3))
    ours = [int(x) for x in results["revise_pick"]]
    print("=" * 78)
    print(f"Revise-and-Pick（3 次生成 + verifier）：准确率 {sum(ours) / len(ours):.3f}")
    print("=" * 78)
    print(f"{'对照':36}{'生成次数':>8}{'准确率':>8}{'Revise-and-Pick 的优势':>24}")
    boot_rng = random.Random(1752)
    for key, name, calls in labels:
        theirs = [int(x) for x in results[key]]
        lo, hi = paired_bootstrap(ours, theirs, boot_rng)
        diff = (sum(ours) - sum(theirs)) / len(ours)
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 36 - width) + f"{calls:8d}{sum(theirs) / len(theirs):8.3f}"
              + f"{diff * 100:+10.1f} 点 [{lo * 100:+.1f}, {hi * 100:+.1f}]")

    print("\n观察：")
    print("  · 和论文选的基线（1 次生成）比，Revise-and-Pick 高了将近 18 个点 —— 看起来是一个大的进步。")
    print("  · 换成计算量相同的基线（都用 3 次生成），优势缩小到 4~8 个点：一半以上的\"提升\"只是因为多用了计算。")
    print("    和最强的同等计算基线（weighted best-of-3）比，区间的下限已经碰到 0。")
    print("  · 消融显示 verifier 没有任何贡献：去掉它，准确率不降反升（区间包含 0）。全部的提升来自\"修改\"")
    print("    （第 2 讲：低温复查能修好手滑）。这篇论文的标题和叙事都该改写 —— 好的是修改，不是挑选。")
    print("  · 一个可信的实验结果至少需要三样东西：计算量匹配的基线、拆开每个组件的消融、配对的置信区间。\n")


if __name__ == "__main__":
    main()