"""
=============================================================================
第 7 讲 · 代码 1：Fisher 信息 —— 为什么"五五开"的问题最有用
=============================================================================

对应讲义：lectures/lecture07_fisher_design.md 第 2、3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
你要估计两个模型的实力差 Δ。每场比较的信息量是多少？什么样的比较最划算？

    演示 1：一场比较的 Fisher 信息 p(1 − p) 随 Δ 怎样变化；模拟验证 Cramér–Rao 下界：
            最大似然估计的方差 ≈ 1 / (n · p(1 − p))。
    演示 2：项目反应理论（IRT）的项目信息函数：题目的难度决定"对谁有用"，区分度决定"有多有用"。
    演示 3：给一个新模型安排对手（"计算机自适应测试"的排行榜版本）：
            随机挑对手 / 总和中等水平的对手比 / 按当前估计挑实力相当的对手，哪个估计得最快？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_fisher_information.py        # 约 3 秒

需要 numpy。依赖本目录的 design_tools.py（它会再导入第 5 讲的 bt_numpy.py）。
"""

import math

import numpy as np

from design_tools import (ABILITY_GRID, bar, bernoulli_information, item_information, log_sigmoid, pad, print_header,
                          sigmoid)


# =============================================================================
# 演示 1
# =============================================================================
def demo_bernoulli(rng, n=400, reps=20000):
    print_header("演示 1：一场比较的信息量 p(1 − p)，以及 Cramér–Rao 下界")
    print(f"真实实力差 Δ，每次实验比 n = {n} 场，用最大似然 Δ̂ = log(胜/负) 估计，重复 {reps} 次。\n")
    print(f"  {pad('Δ', 5)} {pad('胜率 p', 7, 'right')} {pad('信息 p(1−p)', 12, 'right')} {pad('相对五五开', 10, 'right')}  "
          f"{pad('', 20)} {pad('下界 1/(n·I)', 13, 'right')} {pad('实际方差', 9, 'right')} {pad('标准误 0.2 需要的场数', 22, 'right')}")
    for delta in [0.0, 0.5, 1.0, 2.0, 3.0]:
        p = float(sigmoid(delta))
        info = float(bernoulli_information(delta))
        wins = rng.binomial(n, p, size=reps)
        ok = (wins > 0) & (wins < n)                   # 全胜 / 全负时最大似然不存在，这里几乎不会发生
        est = np.log(wins[ok] / (n - wins[ok]))
        crb = 1.0 / (n * info)
        need = math.ceil(1.0 / (info * 0.2 ** 2))
        print(f"  {delta:>5.1f} {p:>7.3f} {info:>12.4f} {info / 0.25:>10.2f}  {bar(info / 0.25, 20)} "
              f"{crb:>13.5f} {est.var():>9.5f} {need:>22d}")
    print("""
观察：
  · 信息 p(1 − p) 在 Δ = 0（五五开）时最大，为 0.25；Δ = 3（胜率 95%）时只剩约五分之一。
    直觉：结果几乎确定的比较，你在比赛之前就知道答案了，比完也没学到什么。
  · Cramér–Rao 下界：任何无偏估计的方差 ≥ 1/(n · I)。最大似然的实际方差几乎正好等于这个下界
    （最大似然是“渐近有效”的，第 5 讲第 2.1 节）。
  · 最后一列：要把 Δ 估到标准误 0.2，五五开的比较要 100 场，Δ = 3 的比较要五倍多。
    同样的标注预算，问“难分高下”的问题，能学到多得多的东西。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_item_information():
    print_header("演示 2：项目信息函数 —— 一道题对哪些人有用？")
    items = [("容易题 a=1, b=−1.5", 1.0, -1.5), ("中等题 a=1, b=0", 1.0, 0.0), ("难题 a=1, b=1.5", 1.0, 1.5),
             ("区分度低 a=0.5, b=0", 0.5, 0.0), ("区分度高 a=2, b=0", 2.0, 0.0)]
    thetas = [-3, -2, -1, 0, 1, 2, 3]
    print("2PL 模型 P(答对) = σ(a(θ − b))。每道题对不同能力 θ 的人的信息 a² p(1 − p)：\n")
    print(f"  {pad('题目', 22)}" + "".join(pad(f"θ={t:+d}", 8, "right") for t in thetas) + pad("最大值", 8, "right"))
    for name, a, b in items:
        row = "".join(f"{item_information(t, a, b):>8.3f}" for t in thetas)
        print(f"  {pad(name, 22)}{row}{a * a / 4:>8.3f}")
    print("""
观察：
  · 每道题的信息在 θ = b（能力等于难度，答对的概率正好一半）时最大，最大值是 a²/4。
  · 难题对高手有用，对新手几乎没用（新手反正答不对）；容易题反过来。所以给所有人发同一套题很浪费：
    对每个人来说，大部分题不是太难就是太容易。
  · 区分度高的题（a = 2）在 θ = b 附近信息很多（1.0，是 a = 1 的 4 倍），但峰很窄；区分度低的题哪里都没什么用。
  · 一套试卷对能力 θ 的信息 = 各题信息之和（“测验信息函数”），能力估计的标准误 ≈ 1/√(总信息)。""")


# =============================================================================
# 演示 3
# =============================================================================
def run_matchmaking(rule, s_true, refs, n_games, prior_sd, rng, checkpoints=()):
    """
    模拟给一个新模型安排 n_games 场比赛，返回在 checkpoints 各个场数时的估计值。

    【参数】rule："random"（随机挑对手）、"fixed"（总是和实力 0 的对手比）、"adaptive"（挑最接近当前估计的对手）
    【思路】网格后验逐场更新（第 6 讲演示 3：今天的后验是明天的先验）：赢了实力为 o 的对手，
        对数后验加上 log σ(θ − o)；输了加上 log σ(o − θ)。比每次从头算快得多。
        （和 design_tools.eap_ability 算的是同一个东西，那个函数每次从头算，适合一次性给出很多答案的情形。）
    【返回】dict：场数 → 后验均值
    """
    grid = ABILITY_GRID
    log_post = -0.5 * (grid / prior_sd) ** 2
    est = 0.0
    out = {}
    for t in range(1, n_games + 1):
        if rule == "random":
            o = float(refs[rng.integers(len(refs))])
        elif rule == "fixed":
            o = 0.0
        else:
            o = float(refs[np.argmin(np.abs(refs - est))])
        win = rng.random() < sigmoid(s_true - o)
        log_post = log_post + log_sigmoid(grid - o if win else o - grid)
        w = np.exp(log_post - log_post.max())
        est = float(np.sum(grid * w) / np.sum(w))
        if t in checkpoints:
            out[t] = est
    return out


def demo_matchmaking(rng, reps=400):
    print_header("演示 3：给新模型安排对手 —— 随机 / 固定中等对手 / 自适应")
    refs = np.linspace(-3.0, 3.0, 25)            # 25 个已知实力的参考模型
    prior_sd = 1.5
    budgets = [5, 10, 20, 50, 100]
    rules = ["random", "fixed", "adaptive"]
    print(f"排行榜上有 {len(refs)} 个实力已知的参考模型（−3 到 3 均匀分布）。新模型的真实实力 s* ~ N(0, {prior_sd}²)。")
    print("每场比赛后用网格后验（先验 N(0, 1.5²)）更新估计。三种安排对手的方式：")
    print("  · 随机：每场随机挑一个参考模型")
    print("  · 固定中等：总是和实力 0 的参考模型比")
    print("  · 自适应：挑实力最接近“当前估计”的参考模型（这场比赛的信息量最大）")
    print(f"重复 {reps} 次（每次一个新的 s*），报告估计误差的均方根（RMSE）：\n")
    print(f"  {pad('比赛场数', 9)}" + "".join(pad(n, 12, "right") for n in ["随机", "固定中等", "自适应"]))
    sq = {k: np.zeros((reps, len(budgets))) for k in rules}
    true_all = rng.normal(0, prior_sd, size=reps)
    for r in range(reps):
        for k in rules:
            est = run_matchmaking(k, true_all[r], refs, budgets[-1], prior_sd, rng, checkpoints=budgets)
            sq[k][r] = [(est[n] - true_all[r]) ** 2 for n in budgets]
    rmse = {k: np.sqrt(v.mean(axis=0)) for k, v in sq.items()}
    for b_idx, n in enumerate(budgets):
        print(f"  {n:>9d}" + "".join(f"{rmse[k][b_idx]:>12.3f}" for k in rules))
    extreme = {}
    for k in rules:
        errs = [(run_matchmaking(k, 2.5, refs, 20, prior_sd, rng, checkpoints=(20,))[20] - 2.5) ** 2 for _ in range(reps)]
        extreme[k] = math.sqrt(float(np.mean(errs)))
    print(f"\n  真实实力 2.5 的新模型（很强），比 20 场后的 RMSE：随机 {extreme['random']:.3f}，"
          f"固定中等 {extreme['fixed']:.3f}，自适应 {extreme['adaptive']:.3f}")
    print("""
观察：
  · 自适应地挑“实力相当”的对手，每个场数下误差都最小。100 场时它的 RMSE（约 0.21）比随机挑对手（约 0.29）
    小约 30%；换算成场数，随机挑对手大约要多比一倍的场次才能达到同样的精度（误差 ∝ 1/√场数）。
  · 固定和中等对手比，平均来看比随机还好一点：新模型的实力大多在 0 附近，和 0 比本来就接近“五五开”；
    随机挑对手时，有不少比赛是和 ±3 的模型比，几乎没有信息。
  · 但对很强的新模型（2.5），固定中等对手的比赛几乎总是赢（胜率 92%），每场的信息只有五五开的四分之一多一点，
    它和随机一样差；自适应会很快把对手换成强的参考模型。
  · 自适应的“鸡生蛋”问题：信息量取决于未知的真实实力，只能用当前估计代替。一开始估计不准，
    挑的对手未必最合适；但每场比赛都在改进估计，很快就能挑准。这就是计算机自适应测试（下一个文件）的思路。
  · 【边界】自适应安排对手让数据依赖于之前的结果。最大似然 / 贝叶斯估计仍然有效（似然原理），
    但“每个模型和谁比、比了多少场”不再是随机分配的，依赖随机分配的简单分析（例如直接比较两个模型的原始胜率）
    会有偏差。""")


def main():
    rng = np.random.default_rng(7)
    demo_bernoulli(rng)
    demo_item_information()
    demo_matchmaking(rng)


if __name__ == "__main__":
    main()