# 教材位置：第2章，第2.3节；仅终端输出。
# 单位、变量与形状：x、prob、y 为 (6,)；期望按点数平方计，方差按点数四次方计。
# 学习目标：由离散分布精确计算变换后的期望和方差，并用模拟结果作对照。
# 核心关系：若 X 为公平骰子点数、Y=X²，则 E[Y]=∑x²P(X=x)，
# Var(Y)=∑(x²-E[Y])²P(X=x)；事件 A 的指示变量满足 E[1_A]=P(A)。
# 数据与输出：使用 6 个可能点数及其概率，输出 Y 的理论矩和 50000 次模拟的样本矩，
# 再计算点数为偶数的概率；理论计算不依赖模拟结果。
# 运行方式：安装 NumPy 后执行 python e02_expectation.py，所有结果输出到终端。

import numpy as np

# arange 的终点 7 不包含在内；x 是形状 (6,) 的浮点数组 [1,...,6]。
# 浮点类型便于随后进行加权计算，数组下标与概率数组的位置一一对应。
x = np.arange(1, 7, dtype=float)
# full 创建 6 个相同的概率质量；这些数之和为 1，每个元素表示一个点数的概率。
prob = np.full(6, 1 / 6)
# ** 2 是逐元素平方，得到 Y 的可能值 [1,4,9,16,25,36]。
# 这里先变换可能值再求期望；一般来说 E[X²] 并不等于 (E[X])²。
y = x ** 2
# y * prob 是逐元素乘法；求和实现有限离散分布的期望公式。
mean_y = np.sum(y * prob)
# 先以理论均值为中心，再按概率加权平方偏差，计算的是 Y 的总体方差。
# 由于枚举了完整分布，这里不使用样本方差的 n-1 分母。
var_y = np.sum((y - mean_y) ** 2 * prob)
# x % 2 == 0 标出偶数点数；astype(float) 将 True/False 转成 1.0/0.0。
# 因此它表示事件“X 为偶数”的指示函数，而不是事件“Y 等于 2”。
even = (x % 2 == 0).astype(float)
# 固定种子使这次模拟可复现；伪随机数生成器的状态会随着抽样推进。
rng = np.random.default_rng(11)
# 每个骰子点数先抽取、再平方；sample 的形状为 (50000,)，对应 Y 的独立模拟观测。
sample = rng.integers(1, 7, size=50000) ** 2
# 理论值由完整概率分布得到；模拟值会随随机样本波动，通常只接近而不完全等于理论值。
print("理论期望 E[Y]、方差 Var(Y)：", mean_y, var_y)
# mean 用样本量作分母；var(ddof=1) 用样本量减 1 作分母，得到总体方差的无偏估计。
# 无偏指重复抽样时估计量的期望等于真值，不表示本次估计没有误差。
print("模拟样本的均值、方差：", sample.mean(), sample.var(ddof=1))
# 对指示函数求加权平均，得到偶数事件的理论概率 3/6。
print("偶数点数的概率：", np.sum(even * prob))
# 浮点运算可能有舍入误差，因此用 isclose 检查理论恒等式，而不直接比较浮点数相等。
assert np.isclose(mean_y, 91 / 6)
assert np.isclose(var_y, 5369 / 36)
assert np.isclose(np.sum(even * prob), 0.5)