# 教材位置：第3章，第3.4节；图 3-3。
# 单位、变量与形状：total 为 (6,)，图中总次数已把失败数平移 r；次数按试验次计，概率无量
# 纲。
# 学习目标：理解离散等待时间的无记忆性，并区分负二项分布记录的失败数和总试验次数。
# 核心关系：独立伯努利试验的成功概率为 p，首次成功所需次数 T∈{1,2,...}；
# P(T>s+t|T>s)=P(T>t)。若等到第 r 次成功前有 F 次失败，则总次数 N=F+r，
# E[N]=r/p，Var(N)=Var(F)=r(1-p)/p²。
# 数据与输出：取 p=0.3、r=3，验证两组无记忆关系，打印若干总次数的点概率及两种计数的均值
# 。
# 配套图件：运行后另保存 e10_waiting.png，教材直接采用该文件。
# 运行方式：安装 NumPy、SciPy、Matplotlib 后执行 python e10_waiting.py；所有输出由理论分
# 布计算。

import numpy as np
from scipy import stats

# p 是每次独立试验成功的概率，r 是等待达到的成功总数；本例中 r 为正整数。
p, r = 0.3, 3
# SciPy 的 geom 计入最后一次成功，最小可能值是 1；它不表示首次成功前的失败次数。
geom = stats.geom(p=p)
# SciPy 的 nbinom(n=r,p=p) 表示第 r 次成功前的失败次数，支持集从 0 开始。
# 参数名 n 在这里代表目标成功次数 r，与别处二项分布中“固定试验总数 n”的含义不同。
failures = stats.nbinom(n=r, p=p)
# s 是已经等待且仍未成功的次数，t 是后续比较的额外等待长度；本例二者均取非负整数。
for s, t in [(3, 2), (5, 4)]:
    # sf(u)=P(T>u)=(1-p)^u；T>s+t 蕴含 T>s，所以条件概率写成两个尾概率的比值。
    # 分母 P(T>s)>0；条件是“前 s 次都未成功”，不能换成“第 s 次已经成功”。
    conditional = geom.sf(s + t) / geom.sf(s)
    # 无记忆性说明剩余等待时间分布和从头等待相同，来自后续伯努利试验与过去独立。
    # isclose 用于数值核对；该关系不是通过这两组例子才成立，而是几何分布的理论性质。
    assert np.isclose(conditional, geom.sf(t))
    print("无记忆性（s、t、条件概率）：", s, t, conditional)
# 至少要做 r 次试验才能获得 r 次成功；这里只展示 r,...,r+5 这六种总次数。
# 上界 r+6 不包含在内，这个展示网格并不包含负二项分布的全部无限支持集。
total = np.arange(r, r + 6)
print("总试验次数：", total)
# 要计算 P(N=total)，先将总次数减去固定成功次数 r，转换为失败数 F。
# 这里 概率质量函数 数组各项为离散点概率；只打印前六项时，它们之和通常小于 1。
print("概率质量函数：", failures.pmf(total - r))
# F 的理论均值是 r(1-p)/p；加上 r 次成功，得到 N 的理论均值 r/p。
print("平均失败次数、平均总次数：", failures.mean(), failures.mean() + r)
# 加常数使均值平移、方差不变，因此同一方差公式同时适用于 F 和 N。
assert np.isclose(failures.mean() + r, r / p)
assert np.isclose(failures.var(), r * (1 - p) / p ** 2)

# 教材图 3-3只使用原例p=0.3、r=3，分别显示首次成功与第r次成功两种等待口径。
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)
# 几何变量从1开始；绘图显示前20个可能编号，理论支持集仍为所有正整数。
geom_k = np.arange(1, 21)
axes[0].bar(geom_k, geom.pmf(geom_k), width=0.75)
axes[0].set(title=f'首次成功：p={p}', xlabel='首次成功所在试验次数',
            ylabel='概率质量')
# 总次数从r开始。输入SciPy时必须减r，因为failures记录的是失败次数。
# 绘图区间比终端列出的3..8更宽，仅改变显示范围，没有改变随机变量的定义。
total_plot = np.arange(r, 31)
axes[1].bar(total_plot, failures.pmf(total_plot - r), width=0.8)
axes[1].axvline(failures.mean() + r, color='C1', linestyle='--',
                label=f'平均总次数 = {failures.mean() + r:g}')
axes[1].set(title=f'第 {r} 次成功：p={p}', xlabel='总试验次数',
            ylabel='概率质量')
axes[1].legend(fontsize=7.8)
fig.savefig('e10_waiting.png', dpi=180, bbox_inches='tight')
plt.close(fig)