"""
=============================================================================
前置讲 P3 · 代码 04：随机性的其他来源 —— 多次采样、示例的选择
=============================================================================

对应讲义：lectures/prereq3_statistics_for_prompts.md §6

【这个脚本要回答的问题】
    除了"抽到了哪些题"，提示实验还有别的随机性来源：
    1. 温度 > 0 时每道题采样好几次，标准误该怎么算？（演示 1：朴素 SE vs 按题聚类的 SE）
    2. 预算固定时，是多出几道题好，还是每道题多采样几次好？（演示 2）
    3. 在真实的小模型上，上面两件事是什么样子？（演示 3：第 1 讲的迷你 GPT）
    4. 同一个提示模板，换几个示例，准确率能差多少？（演示 4：示例的选择也是随机性的来源）

【怎么读输出】
    演示 1–2 是模拟（stats_kit.ItemBank 的采样版本）；演示 3–4 是第 1 讲训练出来的真实小模型。

【运行】python3 04_sampling_variance.py     （约 10 秒；需要第 1 讲迷你 GPT 的缓存，没有的话会先训练约 50 秒）
【依赖】PyTorch（演示 3–4）；import 同目录的 stats_kit.py，第 1 讲的 console.py、tiny_gpt.py、pretrain_world.py。
"""

import math
import os
import random
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct, print_header, print_observations, print_table  # noqa: E402
from stats_kit import ItemBank, cluster_se, naive_se, wilson_ci, z_for  # noqa: E402

BANK = ItemBank()
THETA = BANK.theta_for(0.70)


def true_sampled_accuracy(theta, n_mc=200000, seed=0):
    """采样解码时的真实准确率 E[σ(θ + ε − d)]，用大量模拟估计。"""
    rng = random.Random(seed)
    tot = 0.0
    for _ in range(n_mc):
        x = theta + rng.gauss(0, BANK.sd_interaction) - rng.gauss(0, BANK.sd_difficulty)
        tot += 1 / (1 + math.exp(-x))
    return tot / n_mc


def mean(xs):
    return math.fsum(xs) / len(xs)


def demo1_cluster(rng, truth):
    """演示 1：多次采样时的标准误。"""
    print_header("演示 1：每道题采样 5 次——标准误该怎么算？")
    n, k, reps = 100, 5, 2000
    est, nse, cse, cover_n, cover_c = [], [], [], 0, 0
    z = z_for()
    for _ in range(reps):
        outs = BANK.sampled(THETA, BANK.draw_items(n, rng), k, rng)
        m = mean([x for o in outs for x in o])
        a, b = naive_se(outs), cluster_se(outs)
        est.append(m)
        nse.append(a)
        cse.append(b)
        cover_n += abs(m - truth) <= z * a
        cover_c += abs(m - truth) <= z * b
    sd = math.sqrt(math.fsum((e - mean(est)) ** 2 for e in est) / (reps - 1))
    print(f"{n} 道题 × 每道题 {k} 次采样 = {n * k} 个结果；真实准确率 {truth:.3f}；重复 {reps} 次实验")
    print_table(["", "数值"], [
        ["估计值在不同实验之间的真实标准差", f"{sd:.4f}"],
        ["朴素 SE（把 500 个结果当成相互独立）的平均值", f"{mean(nse):.4f}"],
        ["按题聚类的 SE 的平均值", f"{mean(cse):.4f}"],
        ["朴素 95% 区间的覆盖率", pct(cover_n, reps)],
        ["聚类 95% 区间的覆盖率", pct(cover_c, reps)],
    ], aligns=["<", ">"])
    print_observations([
        f"把 500 个结果当成 500 次独立的试验，算出的 SE（{mean(nse):.4f}）比真实的波动（{sd:.4f}）小了一大截，"
        f"区间的覆盖率只有 {pct(cover_n, reps)}——它'自信'得过了头。",
        "原因是同一道题的 5 次采样并不独立：难题 5 次都容易错，简单题 5 次都容易对。真正独立的单位是**题目**。"
        f"先算每道题的平均正确率、再对题目求标准误（聚类 SE），覆盖率回到了 {pct(cover_c, reps)}。",
        "Miller（2024，'Adding Error Bars to Evals'）专门讨论了这件事：LLM 评测里每道题采样多次、或者一篇文章出好几道题时，"
        "都要按'题'（或'文章'）聚类计算标准误。",
    ])


def demo2_budget(rng):
    """演示 2：固定预算下的 n × k 取舍。"""
    print_header("演示 2：一共只能调用 400 次——多出题，还是每道题多采样几次？")
    reps = 2000
    rows = []
    for n, k in [(400, 1), (200, 2), (100, 4), (50, 8), (25, 16)]:
        est = [mean([x for o in BANK.sampled(THETA, BANK.draw_items(n, rng), k, rng) for x in o]) for _ in range(reps)]
        sd = math.sqrt(math.fsum((e - mean(est)) ** 2 for e in est) / (reps - 1))
        rows.append([f"{n} 道题 × {k} 次", f"{sd:.4f}", f"±{100 * z_for() * sd:.1f} 个百分点"])
    print(f"总调用次数固定为 400；每种分配重复 {reps} 次实验，看估计值的波动")
    print_table(["分配方式", "估计值的标准差", "对应的 95% 区间半宽"], rows)
    print_observations([
        "同样 400 次调用，400 道题各采样 1 次最准；把预算换成'更少的题、每道题多采样几次'，误差越来越大。",
        "估计值的方差 ≈ （题目之间的差异）/ n + （同一道题内部的采样噪声）/ (n·k)。多采样只能缩小第二项，"
        "第一项只有**多出题**才能缩小——而它通常是大头。",
        "所以多次采样有它的用处（第 4 讲的 self-consistency 用投票提高单题的准确率；估计'模型在这道题上有多稳定'），"
        "但如果目的只是'把准确率测准'，钱应该花在更多的题上。",
    ])


def demo3_tiny_gpt_sampling(model, tok, probes, generate):
    """演示 3：迷你 GPT 上的多次采样。"""
    print_header("演示 3：真实的小模型——迷你 GPT 在 96 道知识题上、温度 > 0 时各采样 8 次")
    from pretrain_world import question_text
    rows = []
    for temp in [1.0, 1.5]:
        outs = []
        for p in probes:
            prompt = f"<doc> faq : q : {question_text(p.country, p.relation)} a :"
            outs.append([1 if generate(model, tok, prompt, 1, temperature=temp, seed=s).split()[0] == p.answer else 0
                         for s in range(8)])
        per = [sum(o) for o in outs]
        acc = mean([x for o in outs for x in o])
        rows.append([temp, pct(acc, 1), sum(1 for x in per if x == 8), sum(1 for x in per if 0 < x < 8),
                     sum(1 for x in per if x == 0), f"{naive_se(outs):.4f}", f"{cluster_se(outs):.4f}",
                     f"{cluster_se(outs) / naive_se(outs):.2f}"])
    print_table(["温度", "准确率", "8 次全对的题", "时对时错的题", "8 次全错的题", "朴素 SE", "聚类 SE", "聚类 / 朴素"], rows)
    print_observations([
        "真实模型也是这样：大部分题 8 次全对，少数题时对时错——每道题'有多难'差别很大，同一道题的 8 次采样高度相关。",
        f"两种温度下，聚类 SE 都是朴素 SE 的 {rows[0][-1]} 倍左右。用朴素的 SE，你会以为测得比实际精确得多。",
        "温度越高，时对时错的题越多，准确率也越低：温度不只是'随机性'，它会系统性地改变你测到的结果。"
        "比较两个提示时，温度（以及所有采样参数）必须保持一致，并写进实验记录。",
    ])


def demo4_demo_choice(model, tok, probes, top_next):
    """演示 4：示例的选择。"""
    print_header("演示 4：同一个提示模板，换一组示例——准确率能差多少？")
    from pretrain_world import question_text
    seen = [p for p in probes if p.seen]
    rows, summary = [], {}
    for k in [1, 2]:
        accs, n_used = [], 0
        for draw in range(20):
            rng = random.Random(100 + draw)
            demos = rng.sample(seen, k)
            demo_text = " ".join(f"question : {question_text(d.country, d.relation)} answer : {d.answer} ." for d in demos)
            ok = n = 0
            for p in probes:
                if any(d.country == p.country for d in demos):       # 不拿同一个国家的题当示例
                    continue
                prompt = f"<doc> faq : {demo_text} question : {question_text(p.country, p.relation)} answer :"
                ok += top_next(model, tok, prompt, 1)[0][0] == p.answer
                n += 1
            accs.append(ok / n)
            n_used = n
        m = mean(accs)
        sd = math.sqrt(math.fsum((a - m) ** 2 for a in accs) / (len(accs) - 1))
        lo, hi = wilson_ci(round(m * n_used), n_used)
        summary[k] = (min(accs), max(accs), m, sd, (hi - lo) / 2)
        rows.append([f"{k} 个示例", pct(min(accs), 1), pct(m, 1), pct(max(accs), 1), f"{100 * sd:.1f}",
                     f"±{100 * (hi - lo) / 2:.1f}"])
    print("第 1 讲 06 号脚本的'question / answer'格式（模型不认识这两个词）。每次随机选 k 个示例，"
          "整组题都用**同一组**示例；换 20 组示例")
    print_table(["示例数", "20 组里最低", "平均", "最高", "标准差（百分点）", "单次运行的 Wilson 半宽"], rows)
    s1, s2 = summary[1], summary[2]
    print_observations([
        f"只换了示例（模板、题目、模型都不变），1 个示例时准确率从 {s1[0]:.0%} 到 {s1[1]:.0%}，2 个示例时从 {s2[0]:.1%} 到 {s2[1]:.1%}。"
        "Zhao et al.（2021）在 GPT-3 上报告了同样的现象：示例的选择和顺序能让准确率从接近随机变到接近最好。第 2 讲会专门研究它。",
        f"2 个示例时，20 组之间的标准差是 {s2[3] * 100:.1f} 个百分点，比一次运行自己的标准误（约 {s2[4] / z_for() * 100:.1f}，"
        f"即 Wilson 半宽 ±{s2[4] * 100:.1f} 除以 1.96）还大：只用**一组**示例跑一次，报告的区间根本没有包含"
        "'换一组示例会怎样'这部分不确定性，而它比'换一批题'的不确定性还大。",
        "所以示例的选择（还有顺序、措辞）要当作实验里的一个随机因素：用几组不同的示例各跑一遍，报告平均值和波动；"
        "比较两种提示策略时，每种都要在多组示例上比较，而不是各挑一组'看起来不错的'。",
    ])


def main():
    rng = random.Random(40)
    truth = true_sampled_accuracy(THETA)
    demo1_cluster(rng, truth)
    demo2_budget(rng)
    import torch
    torch.set_num_threads(1)
    from pretrain_world import fact_probes
    from tiny_gpt import generate, get_model, top_next
    model, tok, info = get_model("base", verbose=False)
    probes = fact_probes(info["world"])
    demo3_tiny_gpt_sampling(model, tok, probes, generate)
    demo4_demo_choice(model, tok, probes, top_next)


if __name__ == "__main__":
    main()