# 教材位置：第3章，第3.2节；图 3-1。
# 单位、变量与形状：k、manual 为 (21,)；k 按成功次数计，概率质量和累积概率均无量纲。
# 学习目标：从组合计数推导二项分布的概率质量函数，并核对 SciPy 的分布接口和尾概率端点。
# 核心关系：n 次相互独立、成功概率相同为 p 的伯努利试验中，成功次数 X 满足
# P(X=k)=C(n,k)p^k(1-p)^(n-k)，E[X]=np，Var(X)=np(1-p)。
# 数据与输出：使用 n=20、p=0.3，枚举全部 21 个可能成功次数，输出理论均值、方差和三个事件
# 的概率。
# 运行方式：安装 NumPy、SciPy、Matplotlib 后执行 python e07_binomial.py；理论计算不做随
# 机抽样，并保存 e07_binomial.png。

# comb 返回组合数 C(n,k)，表示从 n 个试验位置中选出 k 个成功位置的方式数。
from math import comb
import numpy as np
from scipy import stats

# n 为固定试验次数，p 为每次成功概率；只规定这两个参数还不够，二项模型也要求独立性和同一
# p。
n, p = 20, 0.3
# arange 不包含上界，所以用 n+1 保留 k=n；k 的形状为 (21,)，覆盖完整支持集 0,...,20。
k = np.arange(n + 1)
# 对固定的 k，p^k(1-p)^(n-k) 是一种成功/失败排列的概率，再乘组合数汇总所有排列。
# 列表推导式为每个 j 计算一个概率，最后转成数组；int(j) 将 NumPy 整数显式转成 Python 整
# 数。
manual = np.array([comb(n, int(j)) * p ** j * (1 - p) ** (n - j) for j in k])
# 将参数固定为一个“冻结分布”对象，后面调用 pmf、cdf、sf 时不必再次传入 n、p。
rv = stats.binom(n=n, p=p)
# pmf(k) 是离散点概率 P(X=k)，不是连续分布的密度；数组输入返回形状相同的概率数组。
# 两种实现逐项核对，有助于发现组合数、指数或参数传递错误。
assert np.allclose(manual, rv.pmf(k))
# 完整支持集上的概率之和为 1；使用 isclose 容纳浮点计算的舍入误差。
assert np.isclose(manual.sum(), 1)
# 将成功次数乘以对应概率再求和，验证按定义计算的期望等于 np。
assert np.isclose(np.sum(k * manual), n * p)
# 这是分布自身的均值和方差，并非某一组观测的样本均值、样本方差。
print("理论均值、方差：", rv.mean(), rv.var())
# cdf(4)=P(X<=4)，包含端点 4；sf(7)=P(X>7)=P(X>=8)，所以必须传 7。
# 对离散分布，P(X>=k) 要用 sf(k-1)；直接用 sf(k) 会漏掉 X=k 的概率。
# sf 直接计算右尾，通常也比先计算 cdf 再从 1 中相减更适合很小的尾概率。
print("P(X=6), P(X<=4), P(X>=8):", rv.pmf(6), rv.cdf(4), rv.sf(7))

# 教材图 3-1使用本例唯一的主参数n=20、p=0.3；两种分布表示来自同一个冻结对象rv。
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)
# manual是上面组合公式的结果；折点与柱高使用相同k，显示完整支持集0到n。
axes[0].bar(k, manual, width=0.7, alpha=0.6, label='组合公式')
axes[0].plot(k, rv.pmf(k), 'o', markersize=3, label='SciPy 概率质量函数')
axes[0].set(title=f'二项分布：n={n}, p={p}', xlabel='成功次数 k',
            ylabel='P(X = k)', xlim=(-0.6, n + 0.6))
axes[0].legend(fontsize=7.8)
# 在支持集左右各加一个点，画出累积分布函数的0与1平台；阶梯按右连续方式连接。
cdf_k = np.arange(-1, n + 2)
axes[1].step(cdf_k, rv.cdf(cdf_k), where='post')
axes[1].set(title='同一模型的累积分布函数', xlabel='阈值 k',
            ylabel='P(X <= k)', xlim=(-1, n + 1), ylim=(-0.03, 1.05))
fig.savefig('e07_binomial.png', dpi=180, bbox_inches='tight')
plt.close(fig)