"""
=============================================================================
第 6 讲 · 代码 3：拉普拉斯近似 vs MCMC；用后验回答"谁是第一"
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
拉普拉斯近似只需要一次优化（找 MAP）加一次求逆，比 MCMC 便宜得多。它什么时候够用、什么时候不够？
有了后验以后，能回答哪些最大似然回答不了的问题？

    演示 1：数据多时，拉普拉斯和 MCMC 几乎一样；有一个"只比了 5 场、全胜"的新模型时，
            拉普拉斯把偏斜的后验说成对称的，均值和区间都有偏差。
    演示 2：排行榜的不确定性：从后验样本直接算"每个模型是第一的概率"、"排名的可信区间"。
    演示 3：重复实验 300 次，检查两种方法的 90% 区间实际覆盖了多少（真值从先验里抽）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_laplace_vs_mcmc.py        # 约 15 秒

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

import math

import numpy as np

from bayes_tools import (fit_newton, hessian, laplace, log_posterior, pad, print_header, run_chains, simulate_bt,
                         summarize_chains)

TAU = 1.0
Z90 = 1.6449          # 标准正态的 95% 分位数：均值 ± 1.645 × 标准差 是 90% 区间


def posterior_samples(i, j, y, M, rng, n_samples=4000, n_chains=4, tau=TAU):
    """
    用 MH 采样后验，返回"减去平均值之后"的样本。

    【思路】按拉普拉斯协方差的形状提议（代码 02 演示 4），从 MAP 出发（省去大部分预热），
        每条链丢掉前 1/5。
    【参数】tau：先验标准差
    【返回】(样本数组 (总样本数, M), 最大 R̂, 最小 ESS)
    """
    mean, cov = laplace(i, j, y, M, tau)
    chol = np.linalg.cholesky(cov)
    log_p = lambda s: log_posterior(s, i, j, y, tau)    # noqa: E731
    chains, _ = run_chains(log_p, [mean] * n_chains, 2.38 / math.sqrt(M), n_samples, rng, chol=chol)
    post = chains[:, n_samples // 5:, :]
    post = post - post.mean(axis=2, keepdims=True)
    rh, es = summarize_chains(post)
    return post.reshape(-1, M), float(rh.max()), float(es.min())


def centered_laplace(i, j, y, M, tau=TAU):
    """拉普拉斯近似下"减去平均值之后"的分数的均值和标准差。"""
    mean, cov = laplace(i, j, y, M, tau)
    P = np.eye(M) - 1.0 / M
    return mean - mean.mean(), np.sqrt(np.diag(P @ cov @ P.T))


# =============================================================================
# 演示 1
# =============================================================================
def demo_compare(rng):
    print_header("演示 1：拉普拉斯近似 vs MCMC —— 数据多时一样，后验偏斜时不一样")
    M = 6
    tau = 2.0           # 这个演示用一个比较宽的先验（"分数大概在 ±4 以内"），后验的偏斜更明显
    true_s = np.array([1.0, 0.6, 0.2, -0.2, -0.6, -1.0])
    print(f"先验 N(0, {tau:.0f}²)。新模型 0 和老模型 1~5 各比了 1 场，5 场全胜。")
    for title, n_old in [("数据充足：5 个老模型之间 1500 条比较", 1500), ("数据稀少：5 个老模型之间 40 条比较", 40)]:
        # 老模型 1~5 之间互相比较；新模型 0 和每个老模型各比 1 场，全胜
        old = np.arange(1, M)
        a = rng.choice(old, size=n_old)
        b = rng.choice(old, size=n_old)
        keep = a != b
        i0, j0 = a[keep], b[keep]
        _, _, y0 = simulate_bt(true_s, len(i0), rng, pairs=(i0, j0))
        i = np.concatenate([i0, np.zeros(5, dtype=int)])
        j = np.concatenate([j0, old])
        y = np.concatenate([y0, np.ones(5, dtype=int)])
        lap_m, lap_sd = centered_laplace(i, j, y, M, tau)
        post, rh, es = posterior_samples(i, j, y, M, rng, tau=tau)
        q = np.quantile(post, [0.05, 0.95], axis=0)
        print(f"\n【{title}】  MCMC：最大 R̂ {rh:.3f}，最小 ESS {es:.0f}")
        print(f"  {pad('模型', 6)} {pad('拉普拉斯均值', 13, 'right')} {pad('MCMC 均值', 10, 'right')} "
              f"{pad('拉普拉斯标准差', 15, 'right')} {pad('MCMC 标准差', 12, 'right')} "
              f"{pad('拉普拉斯 90% 区间', 19, 'right')} {pad('MCMC 90% 区间', 17, 'right')}")
        for k in range(M):
            lci = f"[{lap_m[k] - Z90 * lap_sd[k]:+.2f}, {lap_m[k] + Z90 * lap_sd[k]:+.2f}]"
            mci = f"[{q[0, k]:+.2f}, {q[1, k]:+.2f}]"
            name = "0（新）" if k == 0 else str(k)
            print(f"  {pad(name, 6)} {lap_m[k]:>13.3f} {post[:, k].mean():>10.3f} {lap_sd[k]:>15.3f} "
                  f"{post[:, k].std():>12.3f} {lci:>19s} {mci:>17s}")
    print("""
观察：
  · 新模型 0：它的后验是偏斜的（和代码 01 的 3-0 一样）。拉普拉斯近似以 MAP 为中心、左右对称；
    MCMC 的均值比 MAP 高约 0.25~0.35，90% 区间整体右移，并且“左短右长”。
    如果要回答“新模型的胜率至少是多少”（区间的下界），拉普拉斯会给出过于悲观的答案。
  · 老模型（1~5）：数据充足时两种方法的标准差几乎一样。它们“减去平均值之后”的均值整体差了约 0.05，
    是因为平均值里包含新模型 —— 新模型的偏差被平均分摊到了每个人头上；老模型**之间**的差值两种方法一样。
  · 数据稀少时，老模型的后验也不那么正态了（例如最弱的模型 5 很少赢），两种方法的差别变大。
  · 实践：拉普拉斯近似便宜（一次优化 + 一次求逆），数据多、后验不偏斜时首选；
    小样本、全胜 / 全负、或者要计算尾部概率时，用 MCMC 检查一下。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_leaderboard(rng):
    print_header("演示 2：排行榜 —— 谁是第一？有多大把握？")
    M = 10
    # 模仿真实排行榜：前几名很接近，后面拉开
    true_s = np.array([1.20, 1.10, 1.05, 0.70, 0.40, 0.10, -0.30, -0.80, -1.30, -2.15])
    i, j, y = simulate_bt(true_s, 400, rng)
    s_mle, _ = fit_newton(i, j, y, M, l2=1e-6)
    post, rh, es = posterior_samples(i, j, y, M, rng)
    # 每个样本里，每个模型的排名（0 表示第一）
    ranks = np.argsort(np.argsort(-post, axis=1), axis=1)
    p_first = (ranks == 0).mean(axis=0)
    p_top3 = (ranks < 3).mean(axis=0)
    exp_rank = ranks.mean(axis=0) + 1
    rank_lo = np.quantile(ranks, 0.05, axis=0) + 1
    rank_hi = np.quantile(ranks, 0.95, axis=0) + 1
    order = np.argsort(-s_mle)
    print(f"{M} 个模型、400 条比较（每个模型约 80 场）。MCMC：最大 R̂ {rh:.3f}，最小 ESS {es:.0f}。按最大似然的分数排序：\n")
    print(f"  {pad('最大似然排名', 13)} {pad('模型', 5)} {pad('真实分数', 9, 'right')} {pad('MLE 分数', 9, 'right')} "
          f"{pad('P(第一)', 8, 'right')} {pad('P(前三)', 8, 'right')} {pad('平均排名', 9, 'right')} "
          f"{pad('排名的 90% 区间', 16, 'right')}")
    for r, k in enumerate(order, start=1):
        rng_txt = f"[{rank_lo[k]:.0f}, {rank_hi[k]:.0f}]"
        print(f"  {r:>13d} {pad(f'M{k}', 5)} {true_s[k]:>9.2f} {s_mle[k] - s_mle.mean():>9.2f} {p_first[k]:>8.2f} "
              f"{p_top3[k]:>8.2f} {exp_rank[k]:>9.1f} {rng_txt:>16s}")
    a, b = order[0], order[1]
    p_ab = float((post[:, a] > post[:, b]).mean())
    print(f"""
  最大似然排第一的 M{a} 比第二的 M{b} 强的后验概率：{p_ab:.2f}

观察：
  · 最大似然只给出一个排名；后验给出“每个模型排第几”的整个分布。最大似然排第一的 M0 只有约 63% 的
    后验概率真的是第一，M2 有约 33% —— 400 条比较不足以分出谁是第一。
  · 真实排第二的 M1（1.10）这次运气不好，最大似然把它排到了第 4；但它排名的 90% 区间 [2, 5]
    仍然包含真实名次。只看最大似然的名次，会以为它明显落后于 M3。
  · 排名的可信区间（例如“第 1~3 名”）比单个名次诚实得多。Chatbot Arena 的排行榜就是这样做的：
    它用自助法给每个模型的分数一个置信区间，再报告“排名的上界”（有多少个模型显著地比它强）。
  · 这些问题（“谁是第一”的概率、排名区间、P(A > B)）用后验样本回答非常直接：
    对每个样本做一次排序，再数一数。用最大似然 + 标准误则需要额外的近似。
  · 后验样本还是决策的工具：第 12 讲的 Thompson 采样就是“从后验里抽一个样本，按它行动”。""")


# =============================================================================
# 演示 3
# =============================================================================
def demo_coverage(rng, reps=300):
    print_header("演示 3：重复实验 —— 90% 区间实际覆盖了多少？")
    M, n = 6, 30
    hits = {"lap": 0, "mcmc": 0, "wald": 0}
    wald_missing = 0
    total = 0
    width = {"lap": 0.0, "mcmc": 0.0}
    for _ in range(reps):
        true_s = rng.normal(0, TAU, size=M)                # 真值从先验里抽
        true_c = true_s - true_s.mean()
        i, j, y = simulate_bt(true_s, n, rng)
        lap_m, lap_sd = centered_laplace(i, j, y, M)
        post, _, _ = posterior_samples(i, j, y, M, rng, n_samples=2500, n_chains=2)
        q = np.quantile(post, [0.05, 0.95], axis=0)
        hits["lap"] += int(np.sum(np.abs(true_c - lap_m) <= Z90 * lap_sd))
        hits["mcmc"] += int(np.sum((q[0] <= true_c) & (true_c <= q[1])))
        width["lap"] += float(np.sum(2 * Z90 * lap_sd))
        width["mcmc"] += float(np.sum(q[1] - q[0]))
        # 最大似然 + 标准误（"频率学派"的区间）：数据太少时最大似然可能不存在
        s_mle, _ = fit_newton(i, j, y, M, l2=0.0)
        if not np.all(np.isfinite(s_mle)) or np.max(np.abs(s_mle)) > 8:
            wald_missing += 1
        else:
            Hm = -hessian(s_mle, i, j, y) * n
            cov = np.linalg.pinv(Hm)
            sd = np.sqrt(np.clip(np.diag(cov), 0, None))
            hits["wald"] += int(np.sum(np.abs(true_c - s_mle) <= Z90 * sd))
        total += M
    valid = (reps - wald_missing) * M
    print(f"每次实验：{M} 个选项的真实分数从先验 N(0, 1) 里抽，模拟 {n} 条比较，对“减去平均值之后”的每个分数")
    print(f"求 90% 区间。重复 {reps} 次（共 {total} 个区间）。\n")
    print(f"  {pad('方法', 34)} {pad('覆盖率', 7, 'right')} {pad('平均宽度', 9, 'right')}")
    print(f"  {pad('MCMC（精确后验的样本）', 34)} {hits['mcmc'] / total:>7.3f} {width['mcmc'] / total:>9.2f}")
    print(f"  {pad('拉普拉斯近似', 34)} {hits['lap'] / total:>7.3f} {width['lap'] / total:>9.2f}")
    print(f"  {pad('最大似然 ± 1.645 标准误（能算时）', 34)} {hits['wald'] / max(valid, 1):>7.3f} {'—':>9s}")
    print(f"""
  最大似然不存在（有选项全胜 / 全负，或比较图不连通）的实验：{wald_missing} / {reps}

观察：
  · 真值确实来自先验时，MCMC 的 90% 区间覆盖率接近 90%（代码 01 演示 4 的多维版本）；
    拉普拉斯近似也很接近 —— 对 BT 这种对数凹的后验，正态近似通常不差。
  · 30 条比较、6 个选项时，约 4 成的实验里最大似然根本不存在（给不出区间）。能算的那些实验里，
    最大似然的区间覆盖率还不错，但这是“挑过的”实验（数据恰好比较均衡的那些），不能代表全部。
    先验在小数据时的价值就在这里：任何数据都能给出一个合理的答案。
  · 【边界】这里的“覆盖率好”依赖于“先验是对的”。先验错了（例如真实差距比先验认为的大得多），
    贝叶斯区间也会系统性地偏（代码 01 演示 4 的 Δ = 3、4）。""")


def main():
    rng = np.random.default_rng(606)
    demo_compare(rng)
    demo_leaderboard(rng)
    demo_coverage(rng)


if __name__ == "__main__":
    main()