# 教材位置：第5章，第5.8节；图 5-3。
# 已知观测标准差时的正态—正态共轭更新；本例完全解析，不随机抽样。
# 学习目标：把先验精度、样本信息和后验预测对应到可核对的计算。
# 观测、μ、μ0、τ0、sigma 使用同一测量单位；方差使用该单位的平方。
# data 形状为 (4,)，区间端点形状为 (2,)，所有密度共用原测量尺度。
# 参数：μ0=50、τ0=10、sigma=8，四次观测为 52、48、55、51。
# 输出后验均值/方差/可信区间及下一次观测的预测区间，保存对应图片。
# 运行：安装 requirements.txt 中依赖后执行 python e36_normal_normal.py。

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from plot_style import configure_chinese_plot

# 假设给定 μ 后四个观测条件独立，同服从 N(μ,sigma²)。
# sigma=8 是已知观测标准差，不用这四个观测重新估计它。
mu0, tau0, sigma = 50.0, 10.0, 8.0
data = np.array([52.0, 48.0, 55.0, 51.0])
n = data.size
sample_mean = data.mean()
# 精度是方差的倒数；独立正态信息在精度尺度上相加。
prior_precision = 1 / tau0**2
data_precision = n / sigma**2
posterior_var = 1 / (prior_precision + data_precision)
posterior_mean = posterior_var * (
    prior_precision * mu0 + data_precision * sample_mean
)
posterior_sd = np.sqrt(posterior_var)
# 后验均值是先验均值与样本均值的加权平均，权重和为 1。
data_weight = data_precision / (prior_precision + data_precision)
assert np.isclose(data_weight, 25 / 29)
assert np.isclose(posterior_mean, 2975 / 58)
assert np.isclose(posterior_var, 400 / 29)
prior = stats.norm(loc=mu0, scale=tau0)
posterior = stats.norm(loc=posterior_mean, scale=posterior_sd)
# 等尾可信区间按后验分布分位数给出，是未知 μ 的后验概率范围。
credible = posterior.ppf([0.025, 0.975])
# 预测还含一次新观测自身的噪声，不能只用 posterior_sd 作预测标准差。
predictive_var = sigma**2 + posterior_var
predictive = stats.norm(loc=posterior_mean, scale=np.sqrt(predictive_var))
predictive_interval = predictive.ppf([0.025, 0.975])
assert predictive_var > posterior_var
print("样本量、样本均值：", n, sample_mean)
print("先验精度、数据精度：", prior_precision, data_precision)
print("后验均值、后验方差：", posterior_mean, posterior_var)
print("后验标准差：", posterior_sd)
print("μ 的 95% 等尾可信区间：", credible)
print("下一次观测的预测均值、方差：", predictive.mean(), predictive_var)
print("下一次观测的 95% 预测区间：", predictive_interval)

# 同一后验在两个面板中回答不同问题：参数 μ 与未来观测。
# 图仅呈现以上解析分布，不再生成一批样本以替代精确计算。
configure_chinese_plot()
fig, axes = plt.subplots(1, 2, figsize=(6.3, 3.0), constrained_layout=True)
grid = np.linspace(15, 85, 700)
axes[0].plot(grid, prior.pdf(grid), label="μ 的先验")
axes[0].plot(grid, posterior.pdf(grid), label="μ 的后验")
axes[0].fill_between(
    grid, posterior.pdf(grid),
    where=(grid >= credible[0]) & (grid <= credible[1]), alpha=0.2,
)
axes[0].axvline(sample_mean, color="gray", ls=":", label="样本均值")
axes[0].set(title="未知均值的共轭更新", xlabel="未知均值 μ", ylabel="概率密度")
axes[1].plot(grid, predictive.pdf(grid), label="下一次观测的预测密度")
axes[1].plot(grid, posterior.pdf(grid), ls="--", label="μ 的后验密度")
axes[1].fill_between(
    grid, predictive.pdf(grid),
    where=(grid >= predictive_interval[0]) & (grid <= predictive_interval[1]),
    alpha=0.2,
)
axes[1].set(title="预测需保留观测噪声", xlabel="下一次观测值", ylabel="概率密度")
for ax in axes:
    ax.legend(fontsize=7.5)
    ax.grid(alpha=0.2)
fig.savefig("normal_normal.png", dpi=180, bbox_inches="tight")
plt.close(fig)