# 教材位置：第3章，第3.5节；图 3-4。
# 单位、变量与形状：N、K、n 是对象个数；k 是支持集网格，方差按目标个数平方计。
# 学习目标：比较不放回抽样的超几何分布与独立抽样的二项分布，理解有限总体修正。
# 核心关系：总体 N 个对象中有 K 个成功对象，不放回抽取 n 个时，成功数 X 的
# E[X]=nK/N，Var(X)=n(K/N)(1-K/N)(N-n)/(N-1)；最后一项是有限总体修正因子。
# 数据与输出：设 N=50、K=5、n=10，输出超几何支持集、两种模型的方差及至多一次成功的概率。
# 配套图件：运行后另保存 e11_hypergeom.png，教材直接采用该文件。
# 运行方式：安装 NumPy、SciPy、Matplotlib 后执行 python e11_hypergeom.py；本例只做理论计
# 算，无随机样本。

import numpy as np
from scipy import stats

# 这里的成功对象可以理解为带有某种标记的对象；每个大小为 n 的不放回样本等可能。
# 抽样后总体组成会变化，所以各次“是否成功”的结果不相互独立。
N, K, n = 50, 5, 10
# 特别注意 SciPy 参数命名：M 表示总体数，n 表示总体成功数，N 表示抽取数。
# 因此必须显式映射 M=N、n=K、N=n，不能只按照本脚本变量名直接代入。
hyper = stats.hypergeom(M=N, n=K, N=n)
# 二项对照模型把每次试验当成相互独立、成功概率固定为 K/N，等价于有放回抽样的计数。
# 两种模型的均值相同，但一般概率和方差不同；抽样比例不很小时应谨慎使用二项近似。
binomial = stats.binom(n=n, p=K / N)
# 支持集下界为 max(0,n-(N-K))，上界为 min(n,K)，本例是 0,...,5。
# 不放回时最多只能抽出总体已有的 5 个成功对象，而二项模型允许更多成功。
lower, upper = hyper.support()
# 转成整数后建立包含两端的支持集网格，形状为 (6,)，用于完整概率归一化检查。
k = np.arange(int(lower), int(upper) + 1)
# 本例修正因子为 40/49，小于 1；不放回抽样产生的负相关使计数方差较二项模型小。
# 当 n/N 很小时该因子接近 1；当抽完整个总体时，成功数固定为 K、方差为 0。
correction = (N - n) / (N - 1)
print("支持集端点：", lower, upper)
# 两者都是给定模型的总体方差，不是样本方差或方差估计的误差。
print("超几何、二项模型的方差：", hyper.var(), binomial.var())
# cdf(1) 包含端点 1，即 P(X=0)+P(X=1)；二者差异反映抽样机制不同。
print("P(X<=1):", hyper.cdf(1), binomial.cdf(1))
# hyper.pmf(k) 是完整支持集上的点概率，求和应为 1。
assert np.isclose(hyper.pmf(k).sum(), 1)
# 理论方差应等于二项方差乘有限总体修正因子，核对公式与参数映射是否一致。
assert np.isclose(hyper.var(), binomial.var() * correction)

# 教材图 3-4保留本例N=50、K=5、n=10，不另设总体规模或目标比例的演示情景。
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.8), constrained_layout=True)
# 完整二项支持集为0..n；超几何在k>5时概率质量函数为0，保留这些位置能看见支持集差异。
# 此plot_k不替代上面用于检查超几何归一化的支持集数组k。
plot_k = np.arange(n + 1)
axes[0].bar(plot_k - 0.18, hyper.pmf(plot_k), width=0.36, label='超几何分布')
axes[0].bar(plot_k + 0.18, binomial.pmf(plot_k), width=0.36, label='二项分布')
axes[0].set(title=f'N={N}, K={K}, n={n}', xlabel='目标个数 k',
            ylabel='概率质量', xlim=(-0.7, n + 0.7))
axes[0].legend(fontsize=7.8)
# 右图直接使用两个冻结对象的方差，展示不放回抽样的有限总体修正。
variances = np.array([hyper.var(), binomial.var()])
axes[1].bar(['超几何分布', '二项分布'], variances, color=['C0', 'C1'], width=0.55)
for position, value in enumerate(variances):
    axes[1].text(position, value + 0.02, f'{value:.4f}', ha='center', fontsize=8)
axes[1].set(title=f'有限总体修正 = {correction:.4f}', ylabel='方差',
            ylim=(0, 1.05 * variances.max() + 0.1))
fig.savefig('e11_hypergeom.png', dpi=180, bbox_inches='tight')
plt.close(fig)
