# 教材位置：第5章，第5.7节；图 5-2。
# 单位、变量与形状：grid、density 为 (4001,)；成功概率与参数空间无量纲，区间是未知概率的
# 后验区间。
# 学习目标：用 贝塔—二项 共轭更新估计未知成功概率，并用数值积分核对。
# 核心数学关系：p ~ 贝塔分布(alpha,beta)、K|p ~ 二项分布(n,p) 时，
# p|K=k ~ 贝塔分布(alpha+k,beta+n-k)，后验均值为 (alpha+k)/(alpha+beta+n)。
# 数据与输出：先验 贝塔分布(2,2)，20 次条件独立的同概率试验中成功 15 次；
# 打印解析/积分得到的后验均值和等尾 95% 可信区间，并保存 e17_beta_update.png。
# 运行方式：安装支持 np.trapezoid 的 NumPy、SciPy 及 Matplotlib 后，
# 在终端运行 python e17_beta_update.py。

import numpy as np
from scipy import stats

# alpha、beta 为正的先验形状参数；n 是试验数，k 是成功数，失败数为 n-k。
# 贝塔分布 密度含 p**(alpha-1)*(1-p)**(beta-1)，体现先验对成功率的倾向。
alpha, beta, n, k = 2.0, 2.0, 20, 15
# 对称的 贝塔分布(2,2) 以 0.5 为中心，相比均匀先验更偏向中间成功率。
prior = stats.beta(alpha, beta)
# 似然含 p**k*(1-p)**(n-k)，与先验相乘后指数相加，得到 贝塔分布(17,7)。
# 共轭更新计算的是未知参数 p 的分布，而不是未来成功次数的分布。
posterior = stats.beta(alpha + k, beta + n - k)
# 用 4001 个等距点覆盖参数空间 [0,1]；一维数组形状为 (4001,)。
# 这里积分较平滑且端点密度为零，细网格能给出很精确的数值核对。
grid = np.linspace(0, 1, 4001)
# 逐个候选 p 计算“先验密度 × 观测到恰好 k 次成功的似然”。
# binom.pmf 的第三个参数在这里是整条 grid，所以输出与网格形状相同。
# 得到的函数尚未归一化，其关于 p 的积分是观测数据的边际概率。
unnormalized = prior.pdf(grid) * stats.binom.pmf(k, n, grid)
# trapezoid(y,x) 用梯形法近似积分；除以总面积后得到积分约为 1 的后验密度。
# 连续密度必须按网格间距积分，不能直接把高度的和当作精确的积分。
density = unnormalized / np.trapezoid(unnormalized, grid)
# 连续后验的期望是 integral p*f(p|data) dp；这里使用同一网格数值近似。
grid_mean = np.trapezoid(grid * density, grid)
print("解析后验均值、积分后验均值：", posterior.mean(), grid_mean)
# ppf 接收两个左侧累计概率，给出各留 2.5% 后验质量的等尾可信区间。
# 在本先验和模型下，参数落在该区间内的后验概率为 95%。
print("95% 可信区间：", posterior.ppf([0.025, 0.975]))
# 解析均值 17/(17+7)=17/24；后验均值同时融合先验与样本成功率 15/20。
assert np.isclose(posterior.mean(), 17 / 24)
# 积分存在离散化误差，用绝对容差 atol=1e-6 检查数值结果足够接近解析值。
assert np.isclose(grid_mean, posterior.mean(), atol=1e-6)

# 教材图 5-2中的未知成功率用数学符号theta表示；SciPy的binom接口把同一参数称为p。
# 原数组grid的每个元素都是一个候选theta，既不是观测次数k，也不是新的随机样本。
import matplotlib.pyplot as plt
# 在创建第一幅图之前配置可跨系统选择的中文字体。
from plot_style import configure_chinese_plot

configure_chinese_plot()

fig, axes = plt.subplots(1, 2, figsize=(6.3, 2.9), constrained_layout=True)
# 原积分网格在本例端点密度为0，可直接用于画先验及解析后验，不另设隐藏网格参数。
prior_label = f'先验贝塔({alpha:g}, {beta:g})'
posterior_label = f'后验贝塔({alpha + k:g}, {beta + n - k:g})'
axes[0].plot(grid, prior.pdf(grid), label=prior_label)
axes[0].plot(grid, posterior.pdf(grid), color='C1', label=posterior_label)
# 等尾可信区间直接采用本例posterior的分位数，阴影面积约为95%的后验质量。
credible_lower, credible_upper = posterior.ppf([0.025, 0.975])
axes[0].fill_between(grid, posterior.pdf(grid),
                     where=(grid >= credible_lower) & (grid <= credible_upper),
                     color='C1', alpha=0.2, label='95% 可信区间')
axes[0].axvline(posterior.mean(), color='C1', linestyle='--', linewidth=1)
axes[0].set(title=f'共轭更新：n={n}, k={k}',
            xlabel=r'未知成功概率 $\theta$', ylabel='概率密度')
axes[0].legend(fontsize=7.8)
# 右图对照解析共轭后验与原数值积分的归一化密度；二者接近重合是实现核查。
axes[1].plot(grid, posterior.pdf(grid), linewidth=2.5, alpha=0.6,
             label='解析后验')
axes[1].plot(grid, density, '--', color='C2', linewidth=1.2,
             label='归一化网格密度')
axes[1].set(title='数值积分核对',
            xlabel=r'候选 $\theta$（接口参数 p）', ylabel='后验密度')
axes[1].legend(fontsize=7.8)
fig.savefig('e17_beta_update.png', dpi=180, bbox_inches='tight')
plt.close(fig)