"""
=============================================================================
第 6 讲 · 代码 1：后验分布是什么 —— 两个选项的贝叶斯 BT（网格法，精确计算）
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
一个新模型上了排行榜，和老模型比了 3 场，3 场全胜。它比老模型强多少？
    · 最大似然说：Δ̂ = log(3/0) = +∞（第 5 讲练习 1）—— 显然不合理。
    · 贝叶斯说：先承认"在看到数据之前，差距大概在 ±3 以内"（先验），再用数据更新这个看法（后验）。
      得到的不是一个数，而是一整条分布曲线：Δ 可能是多少、每个值的可能性多大。

只有两个选项时，未知参数只有一个 Δ = s_a − s_b，可以把 Δ 的取值范围切成几千个格点，
在每个格点上直接算"先验 × 似然"，再归一化 —— 这就是精确的后验，不需要任何近似。

    演示 1：不同数据量下的后验：MLE、MAP、后验均值、可信区间、"a 更强"的概率、预测下一场。
    演示 2：先验的影响：数据少时先验很重要，数据多时几乎没影响。
    演示 3：逐条更新 = 一次性更新（今天的后验是明天的先验）；后验宽度按 1/√n 缩小。
    演示 4：可信区间"准不准"：真值来自先验时，90% 区间正好覆盖 90%；对某个固定的真值不一定。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_posterior_basics.py        # 不到 1 秒

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

import math

import numpy as np

from bayes_tools import bar, log_sigmoid, pad, print_header, sigmoid

# 网格：Δ 从 −10 到 10，切成 8001 个点（间隔 0.0025）。范围要足够宽，保证后验在两端几乎为 0。
GRID = np.linspace(-10.0, 10.0, 8001)
DX = GRID[1] - GRID[0]
# 先验：两个选项的分数各自 ~ N(0, 1)，差值 Δ = s_a − s_b ~ N(0, 2)，标准差 √2 ≈ 1.41
PRIOR_SD = math.sqrt(2.0)


def grid_posterior(wins, losses, prior_sd=PRIOR_SD):
    """
    在网格上精确计算 Δ 的后验密度。

    【思路】后验 ∝ 先验 × 似然：
            log 后验(Δ) = log N(Δ; 0, prior_sd²) + wins·log σ(Δ) + losses·log σ(−Δ) + 常数
        在每个格点上算出来，减去最大值（防止 exp 溢出），取 exp，再除以"总面积"（Σ × 格点间隔）归一化。
    【参数】wins、losses：a 赢、输的场数（≥ 0 的整数）；prior_sd：先验标准差 > 0
    【返回】长度 len(GRID) 的数组：每个格点上的后验密度（积分为 1）
    【边界】prior_sd 很大（几乎没有先验）且 losses = 0 时，后验的大部分质量会跑到网格右端，
        此时网格不够宽，结果不可信 —— 这正是"没有先验时，全胜的后验没法归一化"的数值体现。
    """
    log_post = -0.5 * (GRID / prior_sd) ** 2 + wins * log_sigmoid(GRID) + losses * log_sigmoid(-GRID)
    dens = np.exp(log_post - log_post.max())
    return dens / (dens.sum() * DX)


def summarize(dens):
    """
    从网格上的后验密度算出各种摘要。

    【返回】dict：
        mean  ：后验均值 E[Δ]
        sd    ：后验标准差
        map   ：后验最高点（MAP）
        lo, hi：95% 可信区间（等尾：两边各留 2.5%）
        p_pos ：P(Δ > 0 | 数据) —— "a 比 b 强"的后验概率
        pred  ：后验预测 P(下一场 a 赢) = E[σ(Δ)] —— 对所有可能的 Δ 做平均，而不是只用一个点估计
    【思路】密度 × 间隔 = 每个格点的概率质量；累加得到累积分布函数（CDF），在 CDF 上找 2.5% 和 97.5% 的位置。
    """
    w = dens * DX
    mean = float(np.sum(GRID * w))
    sd = float(math.sqrt(np.sum((GRID - mean) ** 2 * w)))
    cdf = np.cumsum(w)
    lo = float(GRID[np.searchsorted(cdf, 0.025)])
    hi = float(GRID[np.searchsorted(cdf, 0.975)])
    return dict(mean=mean, sd=sd, map=float(GRID[np.argmax(dens)]), lo=lo, hi=hi,
                p_pos=float(np.sum(w[GRID > 0])), pred=float(np.sum(sigmoid(GRID) * w)))


def interval(dens, level):
    """等尾可信区间：返回 (下界, 上界)，两边各留 (1 − level)/2。"""
    cdf = np.cumsum(dens * DX)
    a = (1 - level) / 2
    return float(GRID[np.searchsorted(cdf, a)]), float(GRID[np.searchsorted(cdf, 1 - a)])


def mle_text(wins, losses):
    """两个选项的最大似然 Δ̂ = log(wins/losses)；全胜 / 全负时是 ±∞，没有数据时不确定。"""
    if wins == 0 and losses == 0:
        return "任意"
    if losses == 0:
        return "+∞"
    if wins == 0:
        return "−∞"
    return f"{math.log(wins / losses):.3f}"


# =============================================================================
# 演示 1
# =============================================================================
def demo_posteriors():
    print_header("演示 1：同一个先验，不同的数据 —— 后验怎样变化")
    print(f"先验：Δ = s_a − s_b ~ N(0, {PRIOR_SD:.2f}²)。网格 {len(GRID)} 个点，精确计算。\n")
    print(f"  {pad('数据（a 胜-负）', 16)} {pad('MLE', 7, 'right')} {pad('MAP', 7, 'right')} {pad('后验均值', 9, 'right')} "
          f"{pad('后验标准差', 11, 'right')} {pad('95% 可信区间', 18, 'right')} {pad('P(Δ>0)', 7, 'right')} "
          f"{pad('预测下一场', 11, 'right')} {pad('σ(MAP)', 7, 'right')}")
    cases = [(0, 0), (1, 0), (3, 0), (7, 3), (70, 30), (700, 300)]
    for w, l in cases:
        st = summarize(grid_posterior(w, l))
        ci = f"[{st['lo']:+.2f}, {st['hi']:+.2f}]"
        print(f"  {pad(f'{w}-{l}', 16)} {pad(mle_text(w, l), 7, 'right')} {st['map']:>7.3f} {st['mean']:>9.3f} "
              f"{st['sd']:>11.3f} {ci:>18s} {st['p_pos']:>7.3f} {st['pred']:>11.3f} {sigmoid(st['map']):>7.3f}")

    print("\n  3-0 的后验长什么样（每行是 Δ 的一个区间，横条长度 ∝ 后验概率）：")
    dens = grid_posterior(3, 0)
    edges = np.arange(-2.0, 5.01, 0.5)
    masses = [float(np.sum(dens[(GRID >= a) & (GRID < b)]) * DX) for a, b in zip(edges[:-1], edges[1:])]
    top = max(masses)
    for (a, b), m in zip(zip(edges[:-1], edges[1:]), masses):
        print(f"    [{a:+.1f}, {b:+.1f})  {bar(m / top, 40)} {m:.3f}")
    print("""
观察：
  · 没有数据（0-0）时，后验 = 先验：均值 0、标准差 1.41，“a 更强”的概率 0.5。
  · 3-0：最大似然是 +∞，贝叶斯给出一个有限的、合理的答案。后验的最高点（MAP）和均值不一样，
    因为分布是**偏斜**的：左边被“已经赢了 3 场”挡住，右边拖着一条长尾（看上面的横条）。
  · “a 更强”的后验概率随数据增加：3-0 时已经约 0.9，70-30 时接近 1。
  · 最后一列对比两种“预测下一场”：贝叶斯的做法是对所有可能的 Δ 求平均 E[σ(Δ)]；
    “插入”做法是只用一个点估计 σ(MAP)。数据少时，平均的结果更保守（离 0.5 更近）——
    它把“我们还不确定 Δ”这件事也算进了预测里。数据多时两者几乎一样。
  · 70-30 与 700-300：后验均值都接近 log(7/3) = 0.847，标准差缩小到约 1/√10。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_prior_sensitivity():
    print_header("演示 2：先验有多重要？—— 取决于数据有多少")
    sds = [0.5, 1.0, PRIOR_SD, 3.0, 10.0]
    print(f"  {pad('先验标准差', 11)}" + "".join(pad(f"{s:.2f}", 25, "right") for s in sds))
    for w, l in [(3, 0), (7, 3), (70, 30)]:
        row = pad(f"{w}-{l}", 11)
        for s in sds:
            st = summarize(grid_posterior(w, l, s))
            row += pad(f"均值 {st['mean']:.2f}，P(Δ>0) {st['p_pos']:.2f}", 25, "right")
        print(f"  {row}")
    print("""
观察：
  · 3-0：先验标准差从 0.5 变到 10，后验均值从 0.3 变到 5 以上 —— 数据太少，结论主要由先验决定。
    先验标准差 10 相当于“差距可能大到 10”（σ(10) ≈ 0.99995），这个先验本身就不合理。
  · 70-30：先验从 1 变到 10，后验均值几乎不变 —— 100 场比赛的信息量远大于先验。
  · 实践：先验要根据领域知识选（例如“排行榜上两个模型的胜率很少超过 95%，即 |Δ| < 3”），
    并且要报告**换一个合理的先验，结论变不变**（敏感性分析）。""")


# =============================================================================
# 演示 3
# =============================================================================
def demo_sequential():
    print_header("演示 3：逐条更新 = 一次性更新；后验宽度按 1/√n 缩小")
    rng = np.random.default_rng(6)
    true_delta = 1.0
    outcomes = (rng.random(400) < sigmoid(true_delta)).astype(int)
    # 逐条更新：从先验出发，每来一条比较，就乘上这条的似然、再归一化
    log_dens = -0.5 * (GRID / PRIOR_SD) ** 2
    print(f"真实 Δ = {true_delta}（a 每场的胜率 {sigmoid(true_delta):.3f}）。逐条加入比较：\n")
    print(f"  {pad('已看到的场数', 13)} {pad('a 赢的场数', 11, 'right')} {pad('后验均值', 9, 'right')} "
          f"{pad('后验标准差', 11, 'right')} {pad('标准差 × √n', 12, 'right')}")
    checkpoints = {1, 2, 5, 10, 25, 50, 100, 200, 400}
    for n, o in enumerate(outcomes, start=1):
        log_dens = log_dens + (log_sigmoid(GRID) if o == 1 else log_sigmoid(-GRID))
        log_dens = log_dens - log_dens.max()        # 每步减去最大值，防止数值越来越小
        if n in checkpoints:
            dens = np.exp(log_dens)
            dens = dens / (dens.sum() * DX)
            st = summarize(dens)
            print(f"  {n:>13d} {int(outcomes[:n].sum()):>11d} {st['mean']:>9.3f} {st['sd']:>11.3f} "
                  f"{st['sd'] * math.sqrt(n):>12.2f}")
    batch = grid_posterior(int(outcomes.sum()), int(len(outcomes) - outcomes.sum()))
    seq = np.exp(log_dens)
    seq = seq / (seq.sum() * DX)
    diff = float(np.max(np.abs(batch - seq)))
    assert diff < 1e-8
    print(f"""
  逐条更新的最终后验与一次性计算的后验，最大差异 {diff:.1e}（只是浮点误差）。

观察：
  · 贝叶斯更新可以一条一条地做：今天的后验就是明天的先验。这正是在线评分系统（TrueSkill、Glicko）
    的思路 —— 和 Elo 一样逐场更新，但同时记录“对每个选手的分数有多不确定”（第 5 讲第 5 节）。
  · 标准差 × √n 逐渐稳定在一个常数附近（约 1/√(p(1−p)) ≈ 2.2）：后验宽度按 1/√n 缩小，
    和第 5 讲最大似然的标准误一样。数据多了，贝叶斯和最大似然的结论趋于一致（Bernstein–von Mises 定理）。""")


# =============================================================================
# 演示 4
# =============================================================================
def demo_coverage():
    print_header("演示 4：90% 可信区间真的会覆盖真值 90% 吗？")
    n = 10
    level = 0.90
    # 关键技巧：n 场比赛里 a 赢 k 场，后验只依赖 k。所以只有 n + 1 种可能的后验，提前算好它们的区间。
    exact_ci, lap_ci = [], []
    for k in range(n + 1):
        dens = grid_posterior(k, n - k)
        exact_ci.append(interval(dens, level))
        # 拉普拉斯近似：在 MAP 处用曲率（负二阶导数）算正态的标准差
        m = float(GRID[np.argmax(dens)])
        curv = n * sigmoid(m) * sigmoid(-m) + 1.0 / PRIOR_SD ** 2
        half = 1.6449 / math.sqrt(curv)
        lap_ci.append((m - half, m + half))

    def coverage(delta, cis):
        """真值为 delta 时，区间覆盖它的概率 = Σ_k P(赢 k 场 | delta) × 1[delta ∈ 区间_k]（精确计算，不用模拟）。"""
        p = sigmoid(delta)
        total = 0.0
        for k in range(n + 1):
            pk = math.comb(n, k) * p ** k * (1 - p) ** (n - k)
            lo, hi = cis[k]
            total += pk * (lo <= delta <= hi)
        return total

    print(f"每次实验比 {n} 场，用先验 N(0, {PRIOR_SD:.2f}²) 算后验，取 90% 等尾可信区间。覆盖率是精确算出来的。\n")
    print(f"  {pad('真实 Δ', 24)} {pad('精确后验区间的覆盖率', 20, 'right')} {pad('拉普拉斯区间的覆盖率', 20, 'right')}")
    for d in [0.0, 0.5, 1.0, 2.0, 3.0, 4.0]:
        print(f"  {pad(f'固定为 {d:.1f}', 24)} {coverage(d, exact_ci):>20.3f} {coverage(d, lap_ci):>20.3f}")
    # 对先验求平均：∫ 覆盖率(Δ) × 先验(Δ) dΔ
    prior = np.exp(-0.5 * (GRID / PRIOR_SD) ** 2)
    prior /= prior.sum()
    sub = slice(None, None, 20)                          # 隔 20 个格点取一个，够精确
    avg_exact = float(sum(coverage(d, exact_ci) * w for d, w in zip(GRID[sub], prior[sub])) / prior[sub].sum())
    avg_lap = float(sum(coverage(d, lap_ci) * w for d, w in zip(GRID[sub], prior[sub])) / prior[sub].sum())
    print(f"  {pad('按先验随机抽取（平均）', 24)} {avg_exact:>20.3f} {avg_lap:>20.3f}")
    print("""
观察：
  · 最后一行：如果真值真的是从先验里抽出来的，精确后验的 90% 区间平均正好覆盖 90% ——
    这是贝叶斯推断的“自洽”：先验对了，后验的概率说法就是校准的（Cook、Gelman、Rubin 2006 用这一点检查采样代码）。
  · 对某个**固定**的真值，覆盖率可以离 90% 很远：Δ 在 0~2 之间时覆盖率接近 90%；Δ = 3（在先验看来
    很少见的大差距）时只有约 62%；Δ = 4 时是 0 —— 只比 10 场，先验把估计往 0 拉得太厉害，
    即使 10 场全胜，区间的上界也够不着 4。这就是“先验带来偏差”的代价，换来的是数据少时不发散、方差更小。
  · 拉普拉斯区间（正态近似）在这里也大致可用（平均 88%），但在个别真值上偏离精确区间（Δ = 0.5 时 95%）。
    后验越偏斜（全胜、全负），正态近似越差；代码 03 会在多维情形下比较。
  · 覆盖率随 Δ 跳跃式变化，是因为 10 场比赛的结果只有 11 种，区间也只有 11 个。""")


def main():
    demo_posteriors()
    demo_prior_sensitivity()
    demo_sequential()
    demo_coverage()


if __name__ == "__main__":
    main()