"""
=============================================================================
第 7 讲 · 代码 2：计算机自适应测试 —— 以及用它给大语言模型做便宜的评测
=============================================================================

对应讲义：lectures/lecture07_fisher_design.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
GRE、GMAT 这类考试会根据你前面答对答错，挑下一道题。为什么？演示 2 已经说明：一道题只对能力和它的
难度差不多的人有信息量。给所有人发同一套题，对每个人来说大部分题都太难或太容易。

    演示 1：题库 300 道题（2PL 模型）。随机出题 / 固定的中等难度试卷 / 自适应出题（每次挑在当前能力估计处
            信息量最大的题），同样的题数，谁测得准？要测到同样的精度，各需要多少题？
    演示 2：把同样的想法用在大语言模型的评测上（tinyBenchmarks、Fluid Benchmarking 的思路）：
            一个 1000 道题的基准测试，新模型只做几十道题，能不能估计出它在全部 1000 道题上的正确率？
            —— 以及一个陷阱：自适应选题之后，不能直接算"答对的比例"。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_adaptive_testing.py        # 约 10 秒

需要 numpy。依赖本目录的 design_tools.py。
"""

import numpy as np

from design_tools import ABILITY_GRID, irt_prob, item_information, log_sigmoid, pad, print_header, sigmoid


class Posterior:
    """
    一个考生的能力 θ 的网格后验，逐题更新（第 6 讲：今天的后验是明天的先验）。

    【属性】log_post：网格上的对数后验（未归一化）
    【方法】update(答对与否, a, b)；mean()；sd()
    """

    def __init__(self, prior_sd=1.0):
        self.log_post = -0.5 * (ABILITY_GRID / prior_sd) ** 2

    def update(self, correct, a, b):
        eta = a * (ABILITY_GRID - b)
        self.log_post = self.log_post + log_sigmoid(eta if correct else -eta)

    def _weights(self):
        w = np.exp(self.log_post - self.log_post.max())
        return w / w.sum()

    def mean(self):
        return float(np.sum(ABILITY_GRID * self._weights()))

    def sd(self):
        w = self._weights()
        m = float(np.sum(ABILITY_GRID * w))
        return float(np.sqrt(np.sum((ABILITY_GRID - m) ** 2 * w)))


def run_test(theta, a, b, rule, n_items, rng, fixed_order=None, checkpoints=(), stop_sd=None):
    """
    给一个能力为 theta 的考生做一次测验。

    【参数】
        a, b       ：题库的区分度、难度（数组）
        rule       ："random"（随机出题）、"fixed"（按 fixed_order 的顺序出题，人人一样）、
                     "adaptive"（每次挑在当前后验均值处信息量最大、且没做过的题）
        n_items    ：最多做多少道题
        checkpoints：在这些题数时记录后验均值
        stop_sd    ：可选；后验标准差低于它就提前结束（自适应考试常用的"停止规则"）
    【返回】(dict 题数 → 后验均值, 实际做了几道题, 做过的题的下标列表, 答对与否的列表)
    """
    post = Posterior()
    used = np.zeros(len(a), dtype=bool)
    out = {}
    items, answers = [], []
    order = rng.permutation(len(a)) if rule == "random" else fixed_order
    for t in range(n_items):
        if rule == "adaptive":
            info = item_information(post.mean(), a, b)
            info[used] = -1.0                              # 做过的题不再出
            k = int(np.argmax(info))
        else:
            k = int(order[t])
        used[k] = True
        correct = rng.random() < irt_prob(theta, a[k], b[k])
        post.update(correct, a[k], b[k])
        items.append(k)
        answers.append(int(correct))
        if t + 1 in checkpoints:
            out[t + 1] = post.mean()
        if stop_sd is not None and post.sd() < stop_sd:
            return out, t + 1, items, answers
    return out, n_items, items, answers


# =============================================================================
# 演示 1
# =============================================================================
def demo_cat(rng, n_people=300):
    print_header("演示 1：随机出题 / 固定试卷 / 自适应出题")
    n_bank = 300
    a = np.exp(rng.normal(0, 0.3, size=n_bank))            # 区分度：中位数 1，大多在 0.5~2 之间
    b = rng.normal(0, 1.2, size=n_bank)                    # 难度
    fixed = np.argsort(np.abs(b))                          # 固定试卷：最接近"中等难度"的题排在前面
    thetas = rng.normal(0, 1, size=n_people)
    lengths = [5, 10, 20, 40]
    rules = ["random", "fixed", "adaptive"]
    names = {"random": "随机出题", "fixed": "固定中等试卷", "adaptive": "自适应出题"}
    print(f"题库 {n_bank} 道题（2PL：区分度中位数 1，难度 ~ N(0, 1.2²)）；{n_people} 个考生，能力 ~ N(0, 1)。")
    print("能力用网格后验的均值估计（先验 N(0, 1)）。报告估计误差的均方根（RMSE）：\n")
    print(f"  {pad('题数', 6)}" + "".join(pad(names[r], 14, "right") for r in rules))
    est = {r: np.zeros((n_people, len(lengths))) for r in rules}
    for p_idx, th in enumerate(thetas):
        for r in rules:
            out, _, _, _ = run_test(th, a, b, r, lengths[-1], rng, fixed_order=fixed, checkpoints=lengths)
            est[r][p_idx] = [out[n] for n in lengths]
    for k, n in enumerate(lengths):
        print(f"  {n:>6d}" + "".join(f"{np.sqrt(np.mean((est[r][:, k] - thetas) ** 2)):>14.3f}" for r in rules))

    # 按能力分组：中间的人 vs 两头的人（40 道题时）
    mid = np.abs(thetas) < 1.0
    print(f"\n  40 道题时按能力分组的 RMSE：")
    print(f"  {pad('', 18)}" + "".join(pad(names[r], 14, "right") for r in rules))
    for label, mask in [("|θ| < 1（中间）", mid), ("|θ| ≥ 1（两头）", ~mid)]:
        print(f"  {pad(label, 18)}" + "".join(f"{np.sqrt(np.mean((est[r][mask, -1] - thetas[mask]) ** 2)):>14.3f}"
                                             for r in rules))

    # 停止规则：后验标准差 < 0.3 时停止，各需要多少题
    print("\n  停止规则“后验标准差 < 0.3 就停”（最多 100 题），平均要做多少题：")
    for r in ["random", "adaptive"]:
        counts = [run_test(th, a, b, r, 100, rng, fixed_order=fixed, stop_sd=0.3)[1] for th in thetas[:150]]
        print(f"    {pad(names[r], 12)}：平均 {np.mean(counts):.1f} 题，最多 {max(counts)} 题")
    print("""
观察：
  · 同样的题数，自适应出题误差最小；要测到同样的精度，自适应需要的题数大约是随机出题的一半。
  · 固定的中等难度试卷对中间的人很好（题目恰好适合他们），对两头的人很差 —— 高手全对、新手全错，
    答案几乎没有信息，估计只能靠先验（往 0 拉）。随机出题对所有人都一般。
  · 自适应出题对所有能力的人都好：它给高手出难题、给新手出容易题，每道题都尽量“五五开”。
  · 停止规则让每个人做到“够准”就停：自适应考试平均只要随机出题约三分之一的题目。
  · 真实的自适应考试还要处理：题目曝光率（好题被所有人做，容易泄露）、内容平衡、多维能力。""")
    return a, b


# =============================================================================
# 演示 2：LLM 评测
# =============================================================================
def jmle_rasch(Y, sweeps=30):
    """
    从"很多模型 × 很多题"的完整作答矩阵里，联合估计每个模型的能力和每道题的难度（Rasch 模型，第 4 讲）。

    【参数】Y：0/1 矩阵，形状 (模型数, 题数)
    【返回】(能力数组, 难度数组)
    【思路】交替地对每个模型（固定难度）、每道题（固定能力）做一步一维牛顿法；全部向量化。
        第 4 讲的附带参数问题：每个参数背后的数据要足够多。这里每个模型做了 1000 道题、每道题有 40 个模型做过，
        偏差很小。能力加一个很弱的先验（N(0, 3²)），防止全对 / 全错的模型发散。
    """
    n_models, n_items = Y.shape
    theta = np.zeros(n_models)
    b = np.zeros(n_items)
    for _ in range(sweeps):
        p = sigmoid(theta[:, None] - b[None, :])
        g = (Y - p).sum(axis=1) - theta / 9.0
        h = (p * (1 - p)).sum(axis=1) + 1.0 / 9.0
        theta = theta + g / h
        p = sigmoid(theta[:, None] - b[None, :])
        g = -(Y - p).sum(axis=0) - b / 9.0
        h = (p * (1 - p)).sum(axis=0) + 1.0 / 9.0
        b = b + g / h
        shift = b.mean()                                   # 固定零点：难度的平均值为 0（只有差值可识别）
        b -= shift
        theta -= shift
    return theta, b


def demo_llm_eval(rng):
    print_header("演示 2：用 IRT 给大语言模型做便宜的评测")
    n_items, n_old, n_new = 1000, 40, 200
    b_true = rng.normal(0, 1.5, size=n_items)
    theta_old = rng.normal(0, 1.2, size=n_old)
    Y_old = (rng.random((n_old, n_items)) < sigmoid(theta_old[:, None] - b_true[None, :])).astype(int)
    _, b_hat = jmle_rasch(Y_old)
    corr = np.corrcoef(b_hat, b_true)[0, 1]
    print(f"一个 {n_items} 道题的基准测试。先用 {n_old} 个“老模型”的完整作答估计每道题的难度（Rasch，交替牛顿法），")
    print(f"估计的难度和真实难度的相关系数 {corr:.3f}。然后对 {n_new} 个新模型，只让它做 k 道题，估计它在全部 {n_items} 道题上的正确率。\n")
    print("  估计方法：")
    print("    A. 随机 k 道题，直接算答对的比例")
    print("    B. 随机 k 道题，用 IRT 估计能力 θ̂，再预测全部题目的正确率 mean_j σ(θ̂ − b̂_j)")
    print("    C. 自适应 k 道题（每次挑信息量最大的），用 IRT 估计能力再预测")
    print("    D. 自适应 k 道题，直接算答对的比例（错误的做法）\n")
    print(f"  {pad('k', 5)}" + "".join(pad(n, 11, "right") for n in ["A", "B", "C", "D"]) + "   （正确率估计的平均绝对误差）")
    theta_new = rng.normal(0, 1.2, size=n_new)
    ones = np.ones(n_items)
    for k in [10, 30, 100]:
        errs = {m: [] for m in "ABCD"}
        for th in theta_new:
            truth_resp = rng.random(n_items) < sigmoid(th - b_true)     # 新模型在全部题目上的作答
            truth = truth_resp.mean()                                    # 真实的全基准正确率
            # 随机 k 道题
            idx = rng.choice(n_items, size=k, replace=False)
            errs["A"].append(abs(truth_resp[idx].mean() - truth))
            post = Posterior(prior_sd=1.5)
            for j in idx:
                post.update(truth_resp[j], 1.0, b_hat[j])
            errs["B"].append(abs(sigmoid(post.mean() - b_hat).mean() - truth))
            # 自适应 k 道题（作答取自同一份"全部作答"，所以和 A、B 比较公平）
            post = Posterior(prior_sd=1.5)
            used = np.zeros(n_items, dtype=bool)
            picked = []
            for _ in range(k):
                info = item_information(post.mean(), ones, b_hat)
                info[used] = -1.0
                j = int(np.argmax(info))
                used[j] = True
                picked.append(j)
                post.update(truth_resp[j], 1.0, b_hat[j])
            errs["C"].append(abs(sigmoid(post.mean() - b_hat).mean() - truth))
            errs["D"].append(abs(truth_resp[picked].mean() - truth))
        print(f"  {k:>5d}" + "".join(f"{np.mean(errs[m]):>11.3f}" for m in "ABCD"))
    print("""
观察：
  · A（随机抽题、直接算比例）是无偏的，但误差按 1/√k 缩小，10 道题时误差很大。
  · B 用同样的题，只是借助“每道题有多难”的知识（来自老模型）：答对一道难题说明的问题比答对一道容易题多。
    误差明显更小。这就是 tinyBenchmarks（Polo 等 2024）的核心思路：用 IRT 从少量题目推断全基准的分数。
  · C 再加上自适应选题，误差最小（Fluid Benchmarking，Hofmann 等 2025）。
  · D 是个陷阱：自适应选题会挑“五五开”的题，所以无论模型强弱，答对的比例都接近 50%。
    **自适应选题之后，必须用模型（IRT）解释答案，不能直接数答对的比例。**
  · 【边界】这些方法假设 IRT 模型是对的：题目难度对所有模型都一样（没有“这个模型恰好背过这道题”），
    能力是一维的。新模型在某类题上有特殊的强项或弱项时，少量题目的估计会有偏差；题目泄露进训练数据（污染）
    会让难度估计失效。""")


def main():
    rng = np.random.default_rng(77)
    demo_cat(rng)
    demo_llm_eval(rng)


if __name__ == "__main__":
    main()