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

# 设置中文字体，以防乱码
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# --- 1. 伯努利分布 (Bernoulli) ---
# 参数：p (成功概率)
bernoulli_dist = stats.bernoulli(p=0.7)
# 生成1000次伯努利试验（抛硬币）
bernoulli_samples = bernoulli_dist.rvs(size=1000)
print("伯努利分布样本均值（应接近0.7）：", bernoulli_samples.mean())

# --- 2. 二项分布 (Binomial) ---
# 参数：n (试验次数), p (单次成功概率)
binom_dist = stats.binom(n=10, p=0.5)
# 计算：抛10次硬币，恰好得到k次正面的概率
k_values = np.arange(0, 11) # 0到10
prob_k = binom_dist.pmf(k_values) # 概率质量函数

# 绘制二项分布的PMF
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.bar(k_values, prob_k, color='skyblue')
plt.title('二项分布 (n=10, p=0.5)\nPMF')
plt.xlabel('k (成功次数)')
plt.ylabel('概率')
plt.grid(axis='y', alpha=0.3)

# --- 3. 泊松分布 (Poisson) ---
# 参数：λ (单位时间/面积内事件发生的平均次数)
poisson_dist = stats.poisson(mu=3) # 假设平均每小时来3个客户
# 计算：一小时内恰好来k个客户的概率
k_poisson = np.arange(0, 11)
prob_p = poisson_dist.pmf(k_poisson)

plt.subplot(1, 2, 2)
plt.bar(k_poisson, prob_p, color='salmon')
plt.title('泊松分布 (λ=3)\nPMF')
plt.xlabel('k (事件发生次数)')
plt.ylabel('概率')
plt.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.show()

# --- 4. 正态分布 (Normal) ---
# 参数：loc (均值 μ), scale (标准差 σ)
normal_dist = stats.norm(loc=0, scale=1) # 标准正态分布
# 生成随机样本
normal_samples = normal_dist.rvs(size=10000)

plt.figure(figsize=(10, 4))
# 绘制样本直方图（近似概率密度）
plt.subplot(1, 2, 1)
plt.hist(normal_samples, bins=50, density=True, alpha=0.7, color='lightgreen', label='样本直方图')
# 绘制理论概率密度函数(PDF)曲线
x = np.linspace(-4, 4, 100)
pdf_curve = normal_dist.pdf(x)
plt.plot(x, pdf_curve, 'r-', lw=2, label='理论PDF')
plt.title('标准正态分布 (μ=0, σ=1)')
plt.xlabel('x')
plt.ylabel('概率密度')
plt.legend()
plt.grid(alpha=0.3)

# 计算一个区间内的概率（例如，x落在[-1,1]内的概率）
prob_interval = normal_dist.cdf(1) - normal_dist.cdf(-1) # 累积分布函数(CDF)
print(f"标准正态分布下，x落在[-1, 1]区间的概率约为: {prob_interval:.4f}")

plt.subplot(1, 2, 2)
# 绘制不同μ和σ的正态分布PDF
mu_sigma = [(0, 0.5), (0, 1), (0, 2), (1, 1)]
for mu, sigma in mu_sigma:
    x = np.linspace(-4, 4, 200)
    pdf = stats.norm.pdf(x, loc=mu, scale=sigma)
    plt.plot(x, pdf, label=f'μ={mu}, σ={sigma}')
plt.title('不同参数的正态分布PDF')
plt.xlabel('x')
plt.ylabel('概率密度')
plt.legend()
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# --- 5. 计算统计量：期望与方差 ---
# 以二项分布为例，理论值：E[X]=np, Var(X)=np(1-p)
print("\n--- 二项分布 (n=10, p=0.5) 的统计量 ---")
print(f"理论期望 (均值): {binom_dist.mean():.2f}")
print(f"理论方差: {binom_dist.var():.2f}")
print(f"1000次模拟的样本均值: {binom_dist.rvs(size=1000).mean():.2f}")
print(f"1000次模拟的样本方差: {binom_dist.rvs(size=1000).var():.2f}")