import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import beta as beta_dist # 导入beta分布用于绘图

# ------------------------------------------------------------------
# 1. 设置 matplotlib 支持中文显示
# ------------------------------------------------------------------
try:
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows/Linux
    plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
except Exception:
    try:
        plt.rcParams['font.sans-serif'] = ['Arial Unicode MS'] # MacOS
        plt.rcParams['axes.unicode_minus'] = False
    except Exception:
        print("未找到中文字体，绘图可能显示异常。请安装'SimHei'或'Arial Unicode MS'字体。")

# ------------------------------------------------------------------
# 2. 设定模型参数 (为了模拟您图中的效果)
# ------------------------------------------------------------------
n = 20      # X的试验总次数
# 更改 alpha 和 beta 以匹配图中 ~0.72 的均值
alpha = 8.0 # Beta分布的先验参数 alpha
beta = 3.0  # Beta分布的先验参数 beta
# 理论均值 E[Y] = alpha / (alpha + beta) = 8 / 11 ≈ 0.727
theoretical_mean = alpha / (alpha + beta)

# MCMC (Gibbs) 抽样参数
N_samples = 1000   # 总抽样量 N (同您图中的 N=1000)
N_burn_in = 200    # 预估的老化期(预热期) N1

print(f"模型参数: n={n}, alpha={alpha}, beta={beta}")
print(f"理论均值 E[Y]: {theoretical_mean:.4f}")
print(f"抽样设置: 总样本 N={N_samples}")

# ------------------------------------------------------------------
# 3. 初始化
# ------------------------------------------------------------------
# 创建数组来存储所有样本
samples_X = np.zeros(N_samples, dtype=int)
samples_Y = np.zeros(N_samples, dtype=float)

# 设定马尔可夫链的初始状态
# 故意设置一个远离均值(0.727)的初始值，以观察收敛
y_t = 0.1 

# ------------------------------------------------------------------
# 4. 运行 Gibbs 抽样
# ------------------------------------------------------------------
print("开始Gibbs抽样...")
np.random.seed(101) # 使用和您图中一样的随机种子

for i in range(N_samples):
    x_t = np.random.binomial(n, y_t)
    y_t = np.random.beta(x_t + alpha, n - x_t + beta)
    samples_Y[i] = y_t
    samples_X[i] = x_t
print("抽样完成。")

# ------------------------------------------------------------------
# 5. 定义并计算逐步平均值 (Ergodic Mean)
# ------------------------------------------------------------------
def calculate_ergodic_mean(y_samples):
    """
    计算逐步平均值 (累积平均值)
    y_k_bar = (1/k) * sum(y_i for i=1 to k)
    
    使用 np.cumsum() 可以高效实现
    """
    n = len(y_samples)
    # 1. 计算累积和 [y1, y1+y2, y1+y2+y3, ...]
    s = np.cumsum(y_samples)
    # 2. 创建 k 数组 [1, 2, 3, ...]
    k_array = np.arange(1, n + 1)
    # 3. 计算 avg[k] = s[k] / k
    return s / k_array

# 计算所有 Y 样本的逐步平均值 (包括预热期)
ergodic_mean_Y = calculate_ergodic_mean(samples_Y)

# ------------------------------------------------------------------
# 6. 绘制逐步平均值图 (实现您图片中的效果)
# ------------------------------------------------------------------
print("正在绘制逐步平均值图...")

plt.figure(figsize=(10, 6))
plt.plot(np.arange(1, N_samples + 1), ergodic_mean_Y, label=r"逐步平均值 $\bar{y}_k$")
plt.axhline(theoretical_mean, color='red', linestyle='--', label=f"理论均值: {theoretical_mean:.4f}")

# 添加一个垂直线来标记我们估计的预热期
plt.axvline(N_burn_in, color='gray', linestyle=':', label=f"估计的预热期 N1 = {N_burn_in}")

plt.title("使用逐步平均值图查看预热期")
plt.xlabel("迭代次数 (k)")
plt.ylabel(r"逐步平均值 $\bar{y}_k$")
plt.legend()
plt.grid(True)
plt.ylim(0, 1) # Y值在0到1之间
plt.show()