"""
=============================================================================
第 14 讲 · 共享模块：模拟"模型 × 题目"的作答，以及评测数据的几个小工具
=============================================================================

对应讲义：lectures/lecture14_data_for_evaluation.md §0.4

-----------------------------------------------------------------------------
【为什么用模拟】
-----------------------------------------------------------------------------
评测数据的很多问题（一个分数的误差有多大、错标的题怎样改变排名、污染让分数虚高多少、小规模实验能不能
选对数据）都需要"知道真相"才能回答：真实的能力是多少？哪些题被污染了？在真实的基准上我们不知道这些。
所以这里用**项目反应理论**（item response theory，IRT）造一个知道全部真相的小世界：

    每个模型 m 有一个能力 θ_m（可以是向量），每道题 i 有难度 b_i、区分度 a_i、k 个选项。
    模型答对第 i 题的概率：
        P(答对) = 1/k + (1 − 1/k) · σ(a_i · (θ_m − b_i))
    1/k 是瞎猜的概率（四选一是 0.25）；σ 是 sigmoid。能力远高于难度 → 几乎必对；远低于 → 只剩瞎猜。

IRT 是心理测量学里给考试打分的标准模型（GRE、托福都用它），近年也被用来分析语言模型的基准
（例如 tinyBenchmarks，第 14 讲 §6.3）。它当然不是真实的语言模型，但"能力 vs 难度"这个结构足以
演示本讲的大部分现象；讲义会说清楚哪些结论能推广。

这个模块还有：
    - 置信区间（朴素的、按题组聚类的）；
    - n-gram 重叠检测（GPT-3 / DCLM 用 13-gram 做去污染）；
    - 一个只靠计数的三元语言模型（演示 Oren et al. 的"可交换性"污染检验）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 eval_lab.py        # 自检：打印几个模型在三个难度的基准上的期望准确率（不到 1 秒）

依赖 numpy 和第 1 讲的 console.py。
"""

import os
import sys

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture01"))
from console import print_header, print_observations, row  # noqa: E402


# =============================================================================
# 1. IRT：答对的概率
# =============================================================================
def sigmoid(x):
    """
    逻辑函数 σ(x) = 1 / (1 + e^{−x})。

    【边界】x 很大或很小时 np.exp 可能溢出；先把 x 截到 [−40, 40]，σ 在这个范围外已经是 0 或 1（双精度下）。
    """
    return 1.0 / (1.0 + np.exp(-np.clip(x, -40, 40)))


def p_correct(theta, a, b, k=4):
    """
    IRT 的答对概率：P = 1/k + (1 − 1/k) · σ(a · (θ − b))。

    【参数】theta：能力，形状 (M,) 或 (M, 1)——M 个模型；
            a、b：区分度和难度，形状 (n,)——n 道题；
            k：选项个数（k = 4 → 瞎猜 25%；k = None 或 0 表示开放式答案，瞎猜 0%）
    【返回】(M, n) 的概率矩阵
    【numpy 小知识】theta[:, None] 把 (M,) 变成 (M, 1)，和 (n,) 相减时自动"广播"成 (M, n)。
    【边界】a = 0 的题对所有模型都一样（σ(0) = 0.5）——这种题没有区分度，只贡献噪声。
    """
    theta = np.asarray(theta, dtype=float).reshape(-1, 1)
    floor = 1.0 / k if k else 0.0
    return floor + (1 - floor) * sigmoid(a * (theta - b))


def make_items(rng, n, diff_mean=0.0, diff_sd=1.0, disc_mean=1.5, disc_sd=0.3, n_groups=None, group_sd=0.0):
    """
    造一个基准：n 道题的区分度 a 和难度 b。可以有"题组"（同一篇阅读材料下的几道题共享一个难度偏移）。

    【参数】diff_mean、diff_sd：难度的均值和标准差（均值越大越难）；
            disc_mean、disc_sd：区分度（截到 ≥ 0.2，太小的题几乎是纯噪声）；
            n_groups、group_sd：把题分成 n_groups 组，每组加一个 N(0, group_sd²) 的共同难度偏移
    【返回】dict：a、b、group（每道题属于哪一组；没有分组时每题自成一组）
    【数据联系】很多基准都有题组：阅读理解（一篇文章几道题）、MMLU 的 57 个学科、SWE-bench 的 12 个仓库。
              同组的题不独立——讲义 §3.3 的"聚类标准误"就是为它准备的。
    """
    a = np.maximum(rng.normal(disc_mean, disc_sd, n), 0.2)
    b = rng.normal(diff_mean, diff_sd, n)
    if n_groups:
        group = rng.integers(0, n_groups, n)
        b = b + rng.normal(0, group_sd, n_groups)[group]
    else:
        group = np.arange(n)
    return {"a": a, "b": b, "group": group}


def answer(rng, theta, items, k=4):
    """
    让 M 个模型各做一遍这些题：返回 (M, n) 的 0/1 矩阵（1 = 答对）。

    【思路】每道题独立地按 p_correct 掷一次硬币。这代表"模型在这道题上的表现带有随机性"：
            采样解码、提示词的细节、训练随机种子……都会让一个能力固定的模型在边界附近的题上时对时错。
    """
    p = p_correct(theta, items["a"], items["b"], k)
    return (rng.random(p.shape) < p).astype(float)


# =============================================================================
# 2. 置信区间
# =============================================================================
def naive_se(x):
    """
    平均值的朴素标准误：sd / √n（把每道题当作独立样本）。

    【参数】x：形状 (n,) 的 0/1 或实数数组
    【边界】n = 1 时样本标准差没有定义；这里用 ddof=1，n < 2 时返回 nan。
    """
    x = np.asarray(x, dtype=float)
    if len(x) < 2:
        return float("nan")
    return float(np.std(x, ddof=1) / np.sqrt(len(x)))


def clustered_se(x, group):
    """
    按题组聚类的标准误（cluster-robust standard error）。

    【思路】同一组的题一起"被抽到"或"不被抽到"，所以独立的单位是"组"而不是"题"。
            把每组的残差加起来：SE² = Σ_g (Σ_{i∈g} (x_i − x̄))² / n² · G/(G−1)。
            组内的题如果完全独立，这个公式和朴素的 SE 差不多；组内正相关时它更大。
    【参数】x：(n,)；group：(n,) 每道题的组号
    【返回】标准误（浮点数）
    【边界】只有一组时没法估计（G − 1 = 0），返回 nan。
    【数据联系】Miller（2024）*Adding Error Bars to Evals* 建议对有题组的基准都报告聚类标准误。
    """
    x = np.asarray(x, dtype=float)
    resid = x - x.mean()
    _, inv = np.unique(group, return_inverse=True)
    G = inv.max() + 1
    if G < 2:
        return float("nan")
    sums = np.bincount(inv, weights=resid, minlength=G)
    return float(np.sqrt(np.sum(sums ** 2) * G / (G - 1)) / len(x))


# =============================================================================
# 3. n-gram 重叠（去污染检测）
# =============================================================================
_P1, _P2 = np.int64(1_000_003), np.int64(998_244_353)


def ngram_codes(tokens, n):
    """
    把一个整数 token 序列的所有 n-gram 编码成 int64 的"指纹"（多项式滚动哈希）。

    【思路】code = Σ_j tok[i + j] · P1^(n−1−j)  (mod P2)。不同的 n-gram 几乎不会撞到同一个指纹
            （撞上的概率约 1/P2 ≈ 10⁻⁹，对本讲的规模可以忽略）。
    【参数】tokens：一维整数数组；n：n-gram 的长度
    【返回】长度 len(tokens) − n + 1 的 int64 数组；序列比 n 短时返回空数组
    【numpy 小知识】每一步都先取模，保证中间结果不超过 P2² ≈ 10¹⁸ < 2⁶³，int64 不会溢出。
    """
    t = np.asarray(tokens, dtype=np.int64)
    L = len(t) - n + 1
    if L <= 0:
        return np.zeros(0, dtype=np.int64)
    code = np.zeros(L, dtype=np.int64)
    for j in range(n):
        code = (code * _P1 + t[j:j + L]) % _P2
    return code


def corpus_ngram_index(docs, n):
    """
    把一个语料（文档的列表）里所有的 n-gram 指纹收集起来、去重、排序——这就是"去污染索引"。

    【返回】排好序的唯一指纹数组（之后用 np.isin 或 np.searchsorted 查询）
    【边界】只在**文档内部**取 n-gram，不跨文档边界（真实的去污染工具也是这样）。
    【数据联系】GPT-3 和 DCLM 对整个训练语料建 13-gram 索引；工程上通常用 Bloom filter
              （第 4 讲）代替精确的集合，以节省内存。
    """
    parts = [ngram_codes(d, n) for d in docs]
    return np.unique(np.concatenate(parts)) if parts else np.zeros(0, dtype=np.int64)


def overlap_fraction(tokens, index, n):
    """
    一道题的 n-gram 里有多大比例出现在语料索引里。

    【返回】0 到 1 的浮点数；题目比 n 短时返回 0（无法检测——短题是 n-gram 去污染的盲区）
    【numpy 小知识】index 已经排好序，用 np.searchsorted 二分查找，每个查询 O(log N)。
            不要用 np.isin(codes, index)：它每次调用都会把两个数组拼起来重新排序，
            索引有上百万个元素时，1000 道题要慢几百倍。
    """
    codes = ngram_codes(tokens, n)
    if len(codes) == 0 or len(index) == 0:
        return 0.0
    pos = np.minimum(np.searchsorted(index, codes), len(index) - 1)
    return float(np.mean(index[pos] == codes))


# =============================================================================
# 4. 一个只靠计数的三元语言模型（插值平滑）
# =============================================================================
class TrigramLM:
    """
    计数语言模型：p(w | u, v) = λ3·c(u,v,w)/c(u,v) + λ2·c(v,w)/c(v) + λ1·(c(w)+1)/(N+V)。

    【思路】只用计数和线性插值（Jelinek–Mercer 平滑），不需要训练。
            它能"背下"训练语料里出现过的三元组——这正是演示"记忆"和"污染"需要的性质。
    【参数】seqs：训练序列（整数数组）的列表；V：词表大小；lambdas：(λ1, λ2, λ3)，和为 1
    【边界】没见过的上下文 c(u, v) = 0 时，三元部分无定义，把它的权重并到二元/一元上
            （代码里用 np.where 处理）。
    【数据联系】第 6 讲 KenLM（n-gram 语言模型）被用来按困惑度过滤 CCNet；
              这里用同样的东西演示"背下了测试集的顺序"。
    """

    def __init__(self, seqs, V, lambdas=(0.1, 0.3, 0.6), base=None):
        """
        【参数】base：另一个 TrigramLM（可选）。给了它，就在它的计数上"追加"seqs 的计数——
                演示里同一份背景语料要和很多份不同的"污染副本"组合，不必每次从头数背景语料。
        """
        self.V = V
        self.l1, self.l2, self.l3 = lambdas
        uni = np.zeros(V) if base is None else base.uni.copy()
        big = [] if base is None else [base._big_raw]
        tri = [] if base is None else [base._tri_raw]
        for s in seqs:
            s = np.asarray(s, dtype=np.int64)
            uni += np.bincount(s, minlength=V)
            big.append(s[:-1] * V + s[1:])
            tri.append((s[:-2] * V + s[1:-1]) * V + s[2:])
        self.uni = uni
        self.N = uni.sum()
        self._big_raw = np.concatenate(big)
        self._tri_raw = np.concatenate(tri)
        self.big_keys, self.big_cnt = np.unique(self._big_raw, return_counts=True)
        self.tri_keys, self.tri_cnt = np.unique(self._tri_raw, return_counts=True)
        # 二元组的"上下文计数"c(u, v) 就是二元组 (u, v) 的计数；c(v) 用一元计数近似
        self.ctx1 = uni

    @staticmethod
    def _lookup(keys, cnt, q):
        """在排好序的 keys 里查 q 的计数；没有就是 0。"""
        pos = np.searchsorted(keys, q)
        pos = np.minimum(pos, len(keys) - 1)
        return np.where(keys[pos] == q, cnt[pos], 0)

    def logprob_positions(self, u, v, w):
        """
        一批位置的 log p(w | u, v)。u、v、w 是等长的整数数组。

        【返回】自然对数概率的数组
        """
        V = self.V
        u, v, w = (np.asarray(x, dtype=np.int64) for x in (u, v, w))
        # 【边界】u、v、w 可以是任意形状（例如 (排列数, 交界数)），只要三者形状相同
        c_uv = self._lookup(self.big_keys, self.big_cnt, u * V + v)
        c_uvw = self._lookup(self.tri_keys, self.tri_cnt, (u * V + v) * V + w)
        c_v = self.ctx1[v]
        c_vw = self._lookup(self.big_keys, self.big_cnt, v * V + w)
        p1 = (self.uni[w] + 1) / (self.N + V)
        p2 = np.where(c_v > 0, c_vw / np.maximum(c_v, 1), 0.0)
        p3 = np.where(c_uv > 0, c_uvw / np.maximum(c_uv, 1), 0.0)
        l3 = np.where(c_uv > 0, self.l3, 0.0)
        l2 = np.where(c_v > 0, self.l2, 0.0)
        l1 = 1.0 - l2 - l3
        return np.log(l1 * p1 + l2 * p2 + l3 * p3)

    def logprob(self, seq):
        """整条序列的对数概率（从第 3 个 token 开始算，前两个当作上下文）。"""
        s = np.asarray(seq, dtype=np.int64)
        return float(np.sum(self.logprob_positions(s[:-2], s[1:-1], s[2:])))


# =============================================================================
# 自检
# =============================================================================
if __name__ == "__main__":
    print_header("eval_lab 自检：IRT 世界里，5 个模型在三个难度的基准上的期望准确率")
    rng = np.random.default_rng(0)
    thetas = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
    benches = {"容易（难度均值 −1）": make_items(rng, 2000, -1.0),
               "中等（难度均值 0）": make_items(rng, 2000, 0.0),
               "困难（难度均值 2）": make_items(rng, 2000, 2.0)}
    w = [22] + [10] * len(thetas)
    print(row(["基准 \\ 模型能力 θ"] + [f"{t:+.0f}" for t in thetas], w))
    exp_acc = {}
    for name, it in benches.items():
        acc = p_correct(thetas, it["a"], it["b"]).mean(axis=1)
        exp_acc[name] = acc
        print(row([name] + [f"{x:.3f}" for x in acc], w))
    hard = exp_acc["困难（难度均值 2）"]
    print_observations([
        f"同一组模型在困难基准上挤在瞎猜（0.25）附近：θ = −2 和 θ = 0 的期望准确率只差 {hard[2] - hard[0]:.3f}，"
        "有限题数带来的随机误差很容易盖过这个差别——讲义 §6 \"小模型在难基准上没有信号\"。",
        "容易的基准上强模型挤在上限附近（饱和）。一个基准只在\"难度和模型能力相当\"的范围里有区分度。",
    ])