"""
=============================================================================
第 2 讲 · 代码 02：敏感性 —— 模板、示例、顺序，各占多少？
=============================================================================

对应讲义：lectures/lecture02_in_context_learning.md §4

【这个脚本要回答的问题】
    Zhao et al.（2021）和 Lu et al.（2022）在 GPT-3 上发现：同一个任务，换模板、换示例、甚至只换示例的**顺序**，
    准确率都可以从接近随机变到接近最好。在我们的两个迷你 GPT 上：
    1. 4 种模板 × 5 组示例 × 每组 24 种顺序 = 480 个提示，准确率的总波动里，模板、示例、顺序各占多少？（演示 1、3）
    2. 顺序为什么会有影响？是不是和"最后一个示例的标签"有关？（演示 2）
    3. 在 480 个提示里挑开发集上最好的那个，它在新题上还那么好吗？（演示 4：前置讲 P3 的赢家诅咒）

【怎么读输出】
    "方差分解"把 480 个准确率的总方差拆成三份：模板之间（换模板）、同一模板内的示例组之间（换示例）、
    同一组示例的不同顺序之间（换顺序）。三份加起来等于总方差（全方差公式）。

【运行】python3 02_sensitivity.py     （约 50 秒；第一次运行还要训练 bursty 模型约 45 秒）
【依赖】PyTorch；import 同目录的 icl_lab.py，第 1 讲的 console.py，前置讲 P3 的 stats_kit.py。
"""

import itertools
import os
import random
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "prereq3"))
from console import print_header, print_observations, print_table  # noqa: E402
from icl_lab import (TEMPLATES, VERBALIZERS, correct_vector, demo_pool, get_base_model,  # noqa: E402
                     get_bursty_model, label_probs, render, sample_demos, test_set)
from stats_kit import mcnemar_exact, wilson_ci  # noqa: E402

# 4 种"提示设计"：(名字, 模板, 标签词)
DESIGNS = [
    ("sentiment · positive/negative", "sentiment", "positive/negative"),
    ("sentiment · great/awful", "sentiment", "great/awful"),
    ("rating · great/awful", "rating", "great/awful"),
    ("qa · positive/negative", "qa", "positive/negative"),
]
N_SETS = 5


def run_grid(model, tok, test, pool):
    """
    把 4 种设计 × 5 组示例 × 24 种顺序全部跑一遍。

    【返回】列表，每个元素是 {"design", "set", "order"（标签顺序如 "PNNP"）, "correct"（300 道题的 0/1 列表），
        "p_first"（模型给"好评那个标签词"的平均概率；测试集里好评占 46%）}
    【思路】每组示例固定为 2 条好评 + 2 条差评（用 seed 0–4 抽），然后用 itertools.permutations 枚举 4! = 24 种排列。
        同一组示例的 24 种排列，内容完全相同，只有顺序不同——这样"顺序的影响"不会和"内容的影响"混在一起。
    【Python 小知识】itertools.permutations(列表) 依次给出所有排列（元组）。
    """
    out = []
    for dname, tname, vname in DESIGNS:
        t, verb = TEMPLATES[tname], VERBALIZERS[vname]
        words = [verb["positive"], verb["negative"]]
        for s in range(N_SETS):
            demos = sample_demos(pool, random.Random(100 + s), pattern="PPNN")
            for perm in itertools.permutations(demos):
                probs = label_probs(model, tok, [render(t, verb, list(perm), x) for x, _, _ in test], words)
                out.append({"design": dname, "set": s, "order": "".join("P" if d[1] == "positive" else "N" for d in perm),
                            "correct": correct_vector(probs, test), "p_first": sum(p[0] for p in probs) / len(probs)})
    return out


def mean(xs):
    return sum(xs) / len(xs)


def pvar(xs):
    """总体方差（除以 n，不是 n − 1）：全方差公式要求各部分用同一种定义才能加起来。"""
    m = mean(xs)
    return sum((x - m) ** 2 for x in xs) / len(xs)


def decompose(grid, idx=None):
    """
    全方差分解：Var(准确率) = Var(模板均值) + E[Var(示例组均值 | 模板)] + E[Var(准确率 | 模板, 示例组)]。

    【参数】grid：run_grid 的结果；idx：只在这些题目上算准确率（None 表示全部 300 道）
    【返回】(总方差, 模板部分, 示例部分, 顺序部分)——单位是"百分点的平方"
    【思路】把 480 个准确率按 (模板, 示例组) 分组：
        顺序部分 = 每个 (模板, 示例组) 内部 24 个准确率的方差，再取平均；
        示例部分 = 每个模板内部 5 个"示例组平均"的方差，再取平均；
        模板部分 = 4 个"模板平均"的方差。
    【边界】这是描述性的分解（没有做显著性检验）；每个格子里的样本数相等，所以三部分严格加起来等于总方差。
    """
    def acc(r):
        c = r["correct"] if idx is None else [r["correct"][i] for i in idx]
        return 100.0 * sum(c) / len(c)
    designs = sorted({r["design"] for r in grid})
    total = pvar([acc(r) for r in grid])
    tmpl_means, within_set, between_set = [], [], []
    for d in designs:
        set_means = []
        for s in range(N_SETS):
            accs = [acc(r) for r in grid if r["design"] == d and r["set"] == s]
            within_set.append(pvar(accs))
            set_means.append(mean(accs))
        between_set.append(pvar(set_means))
        tmpl_means.append(mean(set_means))
    return total, pvar(tmpl_means), mean(between_set), mean(within_set)


def print_decomposition(grid, title):
    """打印每种设计的均值和范围，以及方差分解表。返回分解结果。"""
    rows = []
    for d, _, _ in DESIGNS:
        accs = [100.0 * mean(r["correct"]) for r in grid if r["design"] == d]
        sd = pvar(accs) ** 0.5
        rows.append([d, f"{mean(accs):.1f}%", f"{min(accs):.1f}%", f"{max(accs):.1f}%", f"{sd:.2f}"])
    print(f"\n  {title}：每种设计 120 个提示（5 组示例 × 24 种顺序）")
    print_table(["设计（模板 · 标签词）", "平均", "最低", "最高", "标准差（百分点）"], rows)
    total, t, s, o = decompose(grid)
    print_table(["方差来源", "方差（百分点²）", "占比"],
                [["换模板 / 标签词", f"{t:.1f}", f"{100 * t / total:.1f}%"],
                 ["同一模板内换示例", f"{s:.2f}", f"{100 * s / total:.1f}%"],
                 ["同一组示例换顺序", f"{o:.2f}", f"{100 * o / total:.1f}%"],
                 ["合计", f"{total:.1f}", "100%"]])
    return total, t, s, o


def demo1(grid_base):
    """演示 1：base 模型的方差分解。"""
    print_header("演示 1：base 模型——480 个提示的准确率，波动从哪里来")
    total, t, s, o = print_decomposition(grid_base, "base 模型")
    sg = [100.0 * mean(r["correct"]) for r in grid_base if r["design"].startswith("sentiment · great")]
    print_observations([
        f"换模板 / 标签词贡献了 {100 * t / total:.1f}% 的方差：四种设计的平均准确率从 70% 左右到 90% 左右。",
        f"换示例只贡献了 {100 * s / total:.1f}%，换顺序 {100 * o / total:.1f}%——但这是**平均**的说法："
        f"在 sentiment · great/awful 这个模型最没把握的设计里，同样 4 条示例只换顺序，准确率就在 "
        f"{min(sg):.1f}%–{max(sg):.1f}% 之间（差 {max(sg) - min(sg):.1f} 个百分点）。",
        "所以'这个模型对顺序不敏感'和'这个提示对顺序很敏感'可以同时成立：敏感性是**提示设计的属性**，不只是模型的属性。",
    ])


def demo2(grid_base, grid_bursty):
    """演示 2：顺序的影响是不是来自最后一个示例的标签。"""
    print_header("演示 2：顺序的影响从哪里来——按'最后一个示例的标签'分组")
    rows = []
    gaps = {}
    for name, grid in [("base", grid_base), ("bursty", grid_bursty)]:
        for d, _, _ in DESIGNS:
            sel = [r for r in grid if r["design"] == d]
            lp = [r for r in sel if r["order"][-1] == "P"]
            ln = [r for r in sel if r["order"][-1] == "N"]
            ap = [100.0 * mean(r["correct"]) for r in lp]
            an = [100.0 * mean(r["correct"]) for r in ln]
            within = mean([max(g) - min(g) for g in (ap, an)])
            gaps[(name, d)] = (mean(ap), mean(an), max(ap + an) - min(ap + an), within,
                               mean([r["p_first"] for r in lp]), mean([r["p_first"] for r in ln]))
            rows.append([name, d, f"{mean(ap):.1f}% / {gaps[(name, d)][4]:.2f}", f"{mean(an):.1f}% / {gaps[(name, d)][5]:.2f}",
                         f"{mean(ap) - mean(an):+.1f}", f"{max(ap + an) - min(ap + an):.1f}", f"{within:.1f}"])
    print_table(["模型", "设计", "最后是好评：准确率 / P(好评词)", "最后是差评：准确率 / P(好评词)", "差值",
                 "120 个的极差", "分组后的组内极差"], rows)
    b = gaps[("base", "sentiment · great/awful")]
    u = gaps[("bursty", "sentiment · positive/negative")]
    ug = gaps[("bursty", "sentiment · great/awful")]
    print_observations([
        f"base 模型的 sentiment · great/awful：最后一个示例是差评时准确率 {b[1]:.1f}%，是好评时 {b[0]:.1f}%；"
        f"120 个提示的极差是 {b[2]:.1f} 个百分点，按最后一个标签分组后组内只剩 {b[3]:.1f}——"
        "这个设计的顺序效应几乎全部来自'最后一个示例写的是哪个标签词'。",
        f"base 模型的方向和'近因偏差'的直觉相反：最后一个示例写了 great，P(great) 平均 {b[4]:.2f}；写了 awful，"
        f"P(great) 反而是 {b[5]:.2f}。它的训练语料里同一页评论互相独立，它没有学会'照抄上一个标签'。",
        f"bursty 模型正好反过来：最后一个示例是好评时，P(好评词) 更高（sentiment · positive/negative："
        f"{u[4]:.2f} vs {u[5]:.2f}），准确率 {u[0]:.1f}% vs {u[1]:.1f}%；great/awful 设计也一样（{ug[0]:.1f}% vs {ug[1]:.1f}%）。"
        "它的语料里同一页的评分相关，它学会了'下一个评分像上一个'——这就是近因偏差（03 号脚本演示 2 直接测量）。",
        "近因偏差往'最后一个标签'那边推。推向好评时，bursty 模型对差评本来就很有把握，推不动，代价小；推向差评时，"
        "一部分把握不大的好评被推错了。所以同样大小的偏差，朝两个方向推，对准确率的影响可以很不一样。",
    ])


def demo3(grid_bursty, base_grid):
    """演示 3：bursty 模型的方差分解，并和 base 模型比较每种设计的波动。"""
    print_header("演示 3：bursty 模型（同页评论相关的语料）——同样的 480 个提示")
    total, t, s, o = print_decomposition(grid_bursty, "bursty 模型")
    sd_b = {d: pvar([100.0 * mean(r["correct"]) for r in base_grid if r["design"] == d]) ** 0.5 for d, _, _ in DESIGNS}
    sd_u = {d: pvar([100.0 * mean(r["correct"]) for r in grid_bursty if r["design"] == d]) ** 0.5 for d, _, _ in DESIGNS}
    more = [d for d in sd_b if sd_u[d] > sd_b[d]]
    print_observations([
        f"bursty 模型在 {len(more)} 种设计上比 base 模型**更**敏感（120 个提示的标准差更大：{'、'.join(more)}）。"
        "它更会'读'上下文里的示例，所以示例的顺序对它也更重要。",
        f"它的总方差里，模板仍是最大的来源（{100 * t / total:.1f}%）；示例和顺序加起来 {100 * (s + o) / total:.1f}%。",
        "Zhao et al.（2021）报告 GPT-3 换一组示例、Lu et al.（2022）报告只换示例的顺序，准确率都可以从接近随机变到接近最好。"
        "我们的两个迷你模型里，越是依赖上下文的那个，越接近这幅图——上下文学习能力和对示例的敏感性是同一枚硬币的两面。",
    ])


def demo4(grid_base, n_items):
    """演示 4：在一批提示里挑开发集第一名。"""
    print_header("演示 4：挑开发集上最好的提示——它在新题上还那么好吗？（base 模型；150 道开发题、150 道留出题）")
    rng = random.Random(0)
    idx = list(range(n_items))
    rng.shuffle(idx)
    dev, held = sorted(idx[:n_items // 2]), sorted(idx[n_items // 2:])

    def acc(r, ids):
        return 100.0 * sum(r["correct"][i] for i in ids) / len(ids)
    rows, picked = [], {}
    scopes = [("全部 480 个提示", grid_base),
              ("sentiment · great/awful 的 120 个", [r for r in grid_base if r["design"] == "sentiment · great/awful"]),
              ("sentiment · positive/negative 的 120 个", [r for r in grid_base if r["design"] == "sentiment · positive/negative"])]
    for scope, grid in scopes:
        ranked = sorted(grid, key=lambda r: acc(r, dev))
        best, worst = ranked[-1], ranked[0]
        k = sum(best["correct"][i] for i in held)
        lo, hi = wilson_ci(k, len(held))
        n01, n10, p = mcnemar_exact([worst["correct"][i] for i in held], [best["correct"][i] for i in held])
        picked[scope] = dict(best_dev=acc(best, dev), best_held=acc(best, held), mean_dev=mean([acc(r, dev) for r in grid]),
                             mean_held=mean([acc(r, held) for r in grid]), worst_held=acc(worst, held), p=p)
        rows.append([scope, f"{best['design'].split(' · ')[0]} · {best['design'].split(' · ')[1]} / 第 {best['set']} 组 / {best['order']}",
                     f"{acc(best, dev):.1f}%", f"{acc(best, held):.1f}% [{100 * lo:.1f}, {100 * hi:.1f}]",
                     f"{picked[scope]['mean_dev']:.1f}% / {picked[scope]['mean_held']:.1f}%",
                     f"{acc(worst, held):.1f}%（{n10} : {n01}，p = {p:.3g}）"])
    print_table(["在哪里挑", "开发集第一名", "它的开发集", "它的留出集 [95% CI]", "全体平均：开发 / 留出",
                 "开发集最后一名的留出集（只有它对 : 只有第一名对，McNemar）"], rows)
    a_all, a_sg, a_sp = (picked[s[0]] for s in scopes)
    print_observations([
        f"在全部 480 个里挑：第一名开发集 {a_all['best_dev']:.1f}%、留出集 {a_all['best_held']:.1f}%，没有回落。"
        "设计之间的差别（qa 比 great/awful 好约 20 个百分点）是真的，远大于 150 道题的抽样误差——挑选有用，"
        f"开发集的第一名和最后一名在留出集上也差得很明显（p = {a_all['p']:.3g}）。",
        f"sentiment · great/awful 里挑：第一名开发集 {a_sg['best_dev']:.1f}% → 留出集 {a_sg['best_held']:.1f}%，"
        f"回落了 {a_sg['best_dev'] - a_sg['best_held']:.1f} 个百分点，但仍比这 120 个的平均（{a_sg['mean_held']:.1f}%）好："
        "它挑中了'最后一个示例是差评'那一类（演示 2），这个效应是真的。",
        f"sentiment · positive/negative 里挑：120 个提示的预测几乎一模一样（标准差 0.18），第一名和平均的差别"
        f"（{a_sp['best_dev']:.1f}% vs {a_sp['mean_dev']:.1f}%）小到可以忽略——没有东西可挑，也就没有'诅咒'。",
        "赢家诅咒最严重的情形是：**很多个候选、真实水平差不多、但各自错在不同的题上**（前置讲 P3 §5 的模拟：20 个一样好的版本，"
        "50 道题时第一名虚高 7.8 个百分点）。迷你模型的候选要么差别很大、要么几乎相同，所以这里只看到了一点点回落。"
        "规矩不变：在开发集上挑，在没用来挑选的留出集上报告，比较两个提示用配对检验。",
    ])


def main():
    test, pool = test_set(), demo_pool()
    base, tok = get_base_model()
    bursty, _ = get_bursty_model()
    grid_base = run_grid(base, tok, test, pool)
    grid_bursty = run_grid(bursty, tok, test, pool)
    demo1(grid_base)
    demo2(grid_base, grid_bursty)
    demo3(grid_bursty, grid_base)
    demo4(grid_base, len(test))


if __name__ == "__main__":
    main()