import numpy as np
import matplotlib.pyplot as plt

def variational_inference_demo():
    """变分推断演示：一维高斯后验的近似"""
    np.random.seed(42)
    
    # 真实后验：混合高斯
    # p(z|x) ∝ 0.4 * N(z|-2, 0.5) + 0.6 * N(z|2, 1.0)
    
    z = np.linspace(-6, 6, 500)
    true_posterior = 0.4 * np.exp(-0.5 * ((z + 2) / 0.5)**2) / (0.5 * np.sqrt(2*np.pi)) + \
                     0.6 * np.exp(-0.5 * ((z - 2) / 1.0)**2) / (1.0 * np.sqrt(2*np.pi))
    true_posterior /= np.trapz(true_posterior, z)
    
    # 变分近似：单高斯（最小化KL(q||p)）
    # 最优q是p的"模式覆盖"近似
    q_mean = np.trapz(z * true_posterior, z)
    q_var = np.trapz((z - q_mean)**2 * true_posterior, z)
    q_approx = np.exp(-0.5 * ((z - q_mean) / np.sqrt(q_var))**2) / (np.sqrt(q_var) * np.sqrt(2*np.pi))
    
    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    axes[0].plot(z, true_posterior, 'b-', linewidth=2, label='True posterior')
    axes[0].plot(z, q_approx, 'r--', linewidth=2, label='Variational approx (mean-field)')
    axes[0].set_xlabel('z')
    axes[0].set_ylabel('p(z)')
    axes[0].set_title('Variational Inference: KL(q||p) minimization')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # ELBO的几何直觉
    # KL(q||p) = ∫q log(q/p) dz
    kl = np.trapz(q_approx * np.log((q_approx + 1e-300) / (true_posterior + 1e-300)), z)
    elbo = -kl  # 简化版本（忽略常数项）
    
    axes[1].bar(['KL(q||p)', '-ELBO'], [kl, -elbo], color=['red', 'blue'], alpha=0.7)
    axes[1].set_ylabel('Value')
    axes[1].set_title('ELBO and KL Divergence')
    axes[1].grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig('variational_inference.png', dpi=150)
    plt.close()
    
    print(f"KL散度: {kl:.4f}")
    print(f"变分近似均值: {q_mean:.3f}")
    print(f"变分近似标准差: {np.sqrt(q_var):.3f}")

variational_inference_demo()

def variational_bayesian_linear_regression():
    """变分贝叶斯线性回归"""
    np.random.seed(42)
    
    # 生成数据
    N = 50
    X_vb = np.sort(np.random.uniform(-3, 3, N))
    t_vb = 2 * X_vb + 1 + np.random.randn(N) * 0.5
    
    Phi = np.column_stack([np.ones(N), X_vb])
    D = Phi.shape[1]
    
    # 变分推断
    # 后验近似: q(w) = N(w|m, S), q(alpha) = Gamma(alpha|a, b)
    
    # 初始化
    m = np.zeros(D)
    S = np.eye(D)
    a_alpha = 1.0
    b_alpha = 1.0
    
    # 先验参数
    a0 = 1.0
    b0 = 1.0
    
    elbos = []
    
    for iteration in range(100):
        # 更新 q(w)
        E_alpha = a_alpha / b_alpha
        S_inv = E_alpha * np.eye(D) + Phi.T @ Phi
        S = np.linalg.inv(S_inv)
        m = S @ Phi.T @ t_vb
        
        # 更新 q(alpha)
        a_alpha = a0 + D / 2
        b_alpha = b0 + 0.5 * (m @ m + np.trace(S))
        
        # 计算ELBO
        E_w_S_inv_w = m @ S_inv @ m
        elbo = -0.5 * N * np.log(2 * np.pi) + 0.5 * D * np.log(E_alpha) - \
               0.5 * np.sum((t_vb - Phi @ m)**2) - 0.5 * E_alpha * (m @ m + np.trace(S))
        elbos.append(elbo)
    
    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 左图：回归结果
    X_plot = np.linspace(-4, 4, 100)
    Phi_plot = np.column_stack([np.ones(100), X_plot])
    
    mean_pred = Phi_plot @ m
    # 预测方差（近似）
    var_pred = np.array([Phi_plot[i] @ S @ Phi_plot[i] for i in range(100)])
    std_pred = np.sqrt(var_pred)
    
    axes[0].scatter(X_vb, t_vb, c='red', s=30, zorder=5, label='Data')
    axes[0].plot(X_plot, mean_pred, 'b-', linewidth=2, label='VB prediction')
    axes[0].fill_between(X_plot, mean_pred - 2*std_pred, mean_pred + 2*std_pred, 
                          alpha=0.2, color='blue', label='95% CI')
    axes[0].plot(X_plot, 2*X_plot + 1, 'g--', linewidth=2, label='True function')
    axes[0].set_xlabel('x')
    axes[0].set_ylabel('y')
    axes[0].set_title('Variational Bayesian Linear Regression')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # 右图：ELBO收敛
    axes[1].plot(elbos)
    axes[1].set_xlabel('Iteration')
    axes[1].set_ylabel('ELBO')
    axes[1].set_title('ELBO Convergence')
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('variational_bayesian_lr.png', dpi=150)
    plt.close()
    
    print(f"变分后验均值: w = {np.round(m, 3)}")
    print(f"真实参数: w = [1, 2]")
    print(f"最终ELBO: {elbos[-1]:.2f}")

variational_bayesian_linear_regression()

def metropolis_hastings():
    """Metropolis-Hastings采样"""
    np.random.seed(42)
    
    # 目标分布：双峰分布
    def log_target(z):
        p1 = np.exp(-0.5 * ((z + 3) / 0.8)**2) / (0.8 * np.sqrt(2*np.pi))
        p2 = np.exp(-0.5 * ((z - 3) / 1.2)**2) / (1.2 * np.sqrt(2*np.pi))
        return np.log(0.4 * p1 + 0.6 * p2 + 1e-300)
    
    # MH采样
    n_samples = 5000
    burn_in = 1000
    samples = np.zeros(n_samples + burn_in)
    z_current = 0.0
    
    accepted = 0
    for i in range(1, n_samples + burn_in):
        # 提议分布：高斯随机游走
        z_proposal = z_current + np.random.randn() * 1.0
        
        # 接受概率
        log_alpha = log_target(z_proposal) - log_target(z_current)
        
        if np.log(np.random.random()) < log_alpha:
            z_current = z_proposal
            accepted += 1
        
        samples[i] = z_current
    
    # 去除burn-in
    samples = samples[burn_in:]
    
    # 可视化
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    # 左图：采样轨迹
    axes[0].plot(samples[:200])
    axes[0].set_xlabel('Iteration')
    axes[0].set_ylabel('Sample')
    axes[0].set_title(f'MH Sampling Trace (acceptance rate: {accepted/(n_samples+burn_in):.2%})')
    axes[0].grid(True, alpha=0.3)
    
    # 中图：样本直方图
    axes[1].hist(samples, bins=50, density=True, alpha=0.7, label='MH samples')
    z_grid = np.linspace(-8, 8, 200)
    target_pdf = np.exp(np.array([log_target(z) for z in z_grid]))
    axes[1].plot(z_grid, target_pdf, 'r-', linewidth=2, label='Target distribution')
    axes[1].set_xlabel('z')
    axes[1].set_ylabel('Density')
    axes[1].set_title('Sample Distribution')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    # 右图：自相关
    max_lag = 50
    autocorr = np.array([np.corrcoef(samples[:-lag], samples[lag:])[0, 1] 
                          for lag in range(1, max_lag+1)])
    axes[2].plot(range(1, max_lag+1), autocorr, 'o-')
    axes[2].set_xlabel('Lag')
    axes[2].set_ylabel('Autocorrelation')
    axes[2].set_title('Autocorrelation')
    axes[2].grid(True, alpha=0.3)
    axes[2].axhline(y=0, color='k', linestyle='--', alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('metropolis_hastings.png', dpi=150)
    plt.close()
    
    print(f"接受率: {accepted/(n_samples+burn_in):.2%}")
    print(f"样本均值: {np.mean(samples):.3f}")
    print(f"样本标准差: {np.std(samples):.3f}")

metropolis_hastings()

def gibbs_sampling():
    """Gibbs采样：从二元高斯分布中采样"""
    np.random.seed(42)
    
    # 目标分布：二元高斯
    mu = np.array([0, 0])
    Sigma = np.array([[1, 0.8], [0.8, 1]])
    Sigma_inv = np.linalg.inv(Sigma)
    
    # 条件分布参数
    # p(z1|z2) = N(mu1 + Sigma12/Sigma22 * (z2 - mu2), Sigma11 - Sigma12^2/Sigma22)
    cond_mean_1_given_2 = lambda z2: mu[0] + Sigma[0,1]/Sigma[1,1] * (z2 - mu[1])
    cond_var_1_given_2 = Sigma[0,0] - Sigma[0,1]**2/Sigma[1,1]
    
    cond_mean_2_given_1 = lambda z1: mu[1] + Sigma[1,0]/Sigma[0,0] * (z1 - mu[0])
    cond_var_2_given_1 = Sigma[1,1] - Sigma[1,0]**2/Sigma[0,0]
    
    # Gibbs采样
    n_samples = 3000
    burn_in = 500
    samples = np.zeros((n_samples + burn_in, 2))
    z = np.array([0.0, 0.0])
    
    for i in range(n_samples + burn_in):
        # 更新 z1
        z[0] = np.random.normal(cond_mean_1_given_2(z[1]), np.sqrt(cond_var_1_given_2))
        # 更新 z2
        z[1] = np.random.normal(cond_mean_2_given_1(z[0]), np.sqrt(cond_var_2_given_1))
        samples[i] = z.copy()
    
    samples = samples[burn_in:]
    
    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 左图：采样点
    axes[0].scatter(samples[:, 0], samples[:, 1], alpha=0.1, s=5, label='Gibbs samples')
    
    # 绘制真实等密度椭圆
    eigenvalues, eigenvectors = np.linalg.eigh(Sigma)
    angle = np.degrees(np.arctan2(eigenvectors[1, 0], eigenvectors[0, 0]))
    for n_std in [1, 2, 3]:
        width = 2 * n_std * np.sqrt(eigenvalues[0])
        height = 2 * n_std * np.sqrt(eigenvalues[1])
        ellipse = plt.Ellipse(mu, width, height, angle=angle, fill=False, 
                               linewidth=2-n_std*0.5, color='red')
        axes[0].add_patch(ellipse)
    
    axes[0].set_xlabel('z1')
    axes[0].set_ylabel('z2')
    axes[0].set_title('Gibbs Sampling from Bivariate Gaussian')
    axes[0].set_xlim(-4, 4)
    axes[0].set_ylim(-4, 4)
    axes[0].set_aspect('equal')
    axes[0].grid(True, alpha=0.3)
    
    # 右图：边缘分布
    axes[1].hist(samples[:, 0], bins=40, density=True, alpha=0.5, label='z1 samples')
    z1_grid = np.linspace(-4, 4, 100)
    axes[1].plot(z1_grid, np.exp(-0.5 * z1_grid**2) / np.sqrt(2*np.pi), 'r-', 
                  linewidth=2, label='N(0,1)')
    axes[1].set_xlabel('z1')
    axes[1].set_ylabel('Density')
    axes[1].set_title('Marginal Distribution of z1')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('gibbs_sampling.png', dpi=150)
    plt.close()
    
    print(f"Gibbs采样统计:")
    print(f"  z1: 均值={np.mean(samples[:, 0]):.3f}, 标准差={np.std(samples[:, 0]):.3f}")
    print(f"  z2: 均值={np.mean(samples[:, 1]):.3f}, 标准差={np.std(samples[:, 1]):.3f}")
    print(f"  相关系数: {np.corrcoef(samples.T)[0, 1]:.3f}")

gibbs_sampling()

def compare_vi_mcmc():
    """比较变分推断和MCMC"""
    np.random.seed(42)
    
    # 目标分布：相关二元高斯
    mu = np.array([0, 0])
    Sigma = np.array([[1, 0.9], [0.9, 1]])
    
    # 变分推断（平均场）
    # 平均场假设 q(z1, z2) = q(z1)q(z2)
    # 最优 q(z1) = N(z1 | mu1, Sigma11)
    vi_mean = mu.copy()
    vi_var = np.diag(Sigma)  # 忽略相关性
    
    # MCMC
    n_samples = 5000
    burn_in = 1000
    samples = np.zeros((n_samples + burn_in, 2))
    z = np.array([0.0, 0.0])
    
    cond_var = 1 - 0.9**2
    
    for i in range(n_samples + burn_in):
        z[0] = np.random.normal(0.9 * z[1], np.sqrt(cond_var))
        z[1] = np.random.normal(0.9 * z[0], np.sqrt(cond_var))
        samples[i] = z.copy()
    samples = samples[burn_in:]
    
    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 变分推断
    z1 = np.linspace(-4, 4, 100)
    z2 = np.linspace(-4, 4, 100)
    Z1, Z2 = np.meshgrid(z1, z2)
    
    # 真实分布
    diff = np.stack([Z1 - mu[0], Z2 - mu[1]], axis=-1)
    true_log_pdf = -0.5 * np.einsum('...i,...ij,...j->...', diff, np.linalg.inv(Sigma), diff)
    true_pdf = np.exp(true_log_pdf)
    true_pdf /= true_pdf.sum()
    
    # 变分近似
    vi_pdf = np.exp(-0.5 * (Z1 - vi_mean[0])**2 / vi_var[0]) / np.sqrt(2*np.pi*vi_var[0]) * \
             np.exp(-0.5 * (Z2 - vi_mean[1])**2 / vi_var[1]) / np.sqrt(2*np.pi*vi_var[1])
    
    axes[0].contourf(Z1, Z2, true_pdf, levels=15, cmap='Blues', alpha=0.7)
    axes[0].contour(Z1, Z2, vi_pdf, levels=10, colors='red', linestyles='--')
    axes[0].set_title('VI (red dashed) vs True (blue contours)')
    axes[0].set_xlabel('z1')
    axes[0].set_ylabel('z2')
    axes[0].set_xlim(-4, 4)
    axes[0].set_ylim(-4, 4)
    axes[0].grid(True, alpha=0.3)
    
    # MCMC
    axes[1].scatter(samples[:, 0], samples[:, 1], alpha=0.05, s=3, c='green')
    axes[1].contour(Z1, Z2, true_pdf, levels=10, colors='blue')
    axes[1].set_title('MCMC Samples vs True Contours')
    axes[1].set_xlabel('z1')
    axes[1].set_ylabel('z2')
    axes[1].set_xlim(-4, 4)
    axes[1].set_ylim(-4, 4)
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('vi_vs_mcmc.png', dpi=150)
    plt.close()
    
    # 比较
    print("变分推断 vs MCMC:")
    print(f"  VI方差(z1): {vi_var[0]:.3f} (真实: {Sigma[0,0]:.3f})")
    print(f"  VI方差(z2): {vi_var[1]:.3f} (真实: {Sigma[1,1]:.3f})")
    print(f"  VI忽略了相关性: 真实 ρ={Sigma[0,1]:.3f}")
    print(f"  MCMC估计 ρ={np.corrcoef(samples.T)[0,1]:.3f}")

compare_vi_mcmc()

def variational_gmm():
    """变分高斯混合模型"""
    np.random.seed(42)
    
    # 生成数据
    N = 300
    X_vgmm = np.vstack([
        np.random.multivariate_normal([0, 0], [[0.5, 0], [0, 0.5]], N//3),
        np.random.multivariate_normal([3, 3], [[0.3, 0.1], [0.1, 0.3]], N//3),
        np.random.multivariate_normal([0, 4], [[0.4, -0.1], [-0.1, 0.4]], N//3)
    ])
    
    K = 5  # 使用比真实组件数更多的组件
    
    # 简化的变分GMM（使用sklearn的BayesianGaussianMixture）
    from sklearn.mixture import BayesianGaussianMixture
    
    bgmm = BayesianGaussianMixture(n_components=K, weight_concentration_prior_type='dirichlet_process',
                                     weight_concentration_prior=0.01, max_iter=200, random_state=42)
    bgmm.fit(X_vgmm)
    
    # 有效的组件数
    effective_components = np.sum(bgmm.weights_ > 0.01)
    
    print(f"设定的组件数: {K}")
    print(f"有效的组件数: {effective_components}")
    print(f"组件权重: {np.round(bgmm.weights_, 3)}")
    
    # 可视化
    fig, ax = plt.subplots(figsize=(8, 6))
    labels = bgmm.predict(X_vgmm)
    ax.scatter(X_vgmm[:, 0], X_vgmm[:, 1], c=labels, cmap='viridis', alpha=0.5, s=10)
    ax.scatter(bgmm.means_[:, 0], bgmm.means_[:, 1], c='red', s=200, marker='*', 
               edgecolors='black', zorder=5)
    ax.set_title(f'Variational GMM (effective K={effective_components})')
    ax.set_xlabel('x1')
    ax.set_ylabel('x2')
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('variational_gmm.png', dpi=150)
    plt.close()

variational_gmm()

def mcmc_bayesian_inference():
    """MCMC用于贝叶斯参数推断"""
    np.random.seed(42)
    
    # 数据：泊松分布的观测
    N = 50
    true_lambda = 3.0
    data = np.random.poisson(true_lambda, N)
    
    # 贝叶斯推断：lambda的后验
    # 先验: lambda ~ Gamma(1, 1)
    # 后验: lambda | data ~ Gamma(1 + sum(data), 1 + N)
    
    # MCMC采样
    n_samples = 5000
    burn_in = 1000
    samples = np.zeros(n_samples + burn_in)
    lambda_current = 2.0
    
    for i in range(1, n_samples + burn_in):
        # 提议分布：对数正态随机游走
        lambda_proposal = lambda_current * np.exp(np.random.randn() * 0.2)
        
        # 对数后验
        def log_posterior(lam):
            if lam <= 0:
                return -np.inf
            # 先验: Gamma(1, 1)
            log_prior = 0  # Gamma(1,1) = Exp(1)
            # 似然: Poisson
            log_lik = np.sum(data * np.log(lam) - lam - np.log(np.arange(1, max(data)+2).cumprod()))
            return log_prior + log_lik
        
        log_alpha = log_posterior(lambda_proposal) - log_posterior(lambda_current)
        log_alpha += np.log(lambda_proposal) - np.log(lambda_current)  # Jacobian
        
        if np.log(np.random.random()) < log_alpha:
            lambda_current = lambda_proposal
        
        samples[i] = lambda_current
    
    samples = samples[burn_in:]
    
    # 解析后验
    alpha_post = 1 + np.sum(data)
    beta_post = 1 + N
    
    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 左图：MCMC采样
    axes[0].hist(samples, bins=50, density=True, alpha=0.7, label='MCMC samples')
    lam_grid = np.linspace(0, 8, 200)
    analytical_pdf = lam_grid**(alpha_post-1) * np.exp(-beta_post*lam_grid) * beta_post**alpha_post / np.math.factorial(int(alpha_post)-1) if alpha_post == int(alpha_post) else lam_grid**(alpha_post-1) * np.exp(-beta_post*lam_grid)
    from scipy.stats import gamma as gamma_dist
    analytical_pdf = gamma_dist.pdf(lam_grid, alpha_post, scale=1/beta_post)
    axes[0].plot(lam_grid, analytical_pdf, 'r-', linewidth=2, label='Analytical posterior')
    axes[0].axvline(x=true_lambda, color='g', linestyle='--', label=f'True λ={true_lambda}')
    axes[0].set_xlabel('λ')
    axes[0].set_ylabel('Density')
    axes[0].set_title('Bayesian Inference for Poisson Rate')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # 右图：收敛诊断
    axes[1].plot(np.cumsum(samples) / np.arange(1, len(samples)+1))
    axes[1].axhline(y=true_lambda, color='g', linestyle='--', label='True λ')
    axes[1].axhline(y=np.mean(samples), color='r', linestyle='-', label='MCMC mean')
    axes[1].set_xlabel('Number of samples')
    axes[1].set_ylabel('Running mean')
    axes[1].set_title('Convergence Diagnostic')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('mcmc_bayesian.png', dpi=150)
    plt.close()
    
    print(f"真实 λ: {true_lambda}")
    print(f"MCMC估计: {np.mean(samples):.3f} ± {np.std(samples):.3f}")
    print(f"解析后验均值: {alpha_post/beta_post:.3f}")

mcmc_bayesian_inference()

def model_comparison_waic():
    """使用WAIC进行模型比较"""
    np.random.seed(42)
    
    # 生成数据
    N = 100
    X_mc = np.random.randn(N, 3)
    true_w = np.array([1, -0.5, 0])
    t_mc = X_mc @ true_w + np.random.randn(N) * 0.5
    
    # 模型1：完整模型（3个特征）
    # 模型2：简化模型（2个特征）
    
    from sklearn.linear_model import BayesianRidge
    
    # 完整模型
    br_full = BayesianRidge()
    br_full.fit(X_mc, t_mc)
    
    # 简化模型
    br_simple = BayesianRidge()
    br_simple.fit(X_mc[:, :2], t_mc)
    
    # 计算对数似然
    from scipy.stats import norm
    
    def compute_waic(model, X, t):
        pred_mean, pred_std = model.predict(X, return_std=True)
        log_lik = norm.logpdf(t, loc=pred_mean, scale=pred_std)
        
        # WAIC = -2 * (lppd - p_waic)
        lppd = np.sum(np.log(np.mean(np.exp(log_lik))))  # 简化
        p_waic = np.sum(np.var(log_lik, axis=0))
        
        return -2 * (np.sum(log_lik) - p_waic)
    
    waic_full = compute_waic(br_full, X_mc, t_mc)
    waic_simple = compute_waic(br_simple, X_mc[:, :2], t_mc)
    
    print(f"完整模型 WAIC: {waic_full:.2f}")
    print(f"简化模型 WAIC: {waic_simple:.2f}")
    print(f"ΔWAIC: {abs(waic_full - waic_simple):.2f}")
    
    if waic_full < waic_simple:
        print("→ 完整模型更优")
    else:
        print("→ 简化模型更优（奥卡姆剃刀）")

model_comparison_waic()