## 实现一下重要抽样
import numpy as np
import matplotlib.pyplot as plt
import scipy as stats
import seaborn as sns

np.random.seed(42)

def baysian_linear_regression_inportance_sampling():
    """
    使用贝叶斯线性回归方法进行重要抽样

    本例需要估计一个简单的线性回归模型的后验分布

    y = β0 + β1*x + ε,  ε ~ N(0, σ²)
      = [1 x] * [β0 β1]' + ε
    """

    # 1. 定义关键数据 
    x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
    y = np.array([2.1, 3.9, 6.2, 8.1, 9.8, 12.3, 13.9, 16.2, 17.8, 20.1])
    n = len(x)

    print(f"数据点数量：{n}")
    print(f"x: {x}")
    print(f"y: {y}")

    # 2. 定义先验分布
    # β0 ~ N(0, 10)
    beta0_prior_mean = 0
    beta0_prior_var = 10

    # β1 ~ N(1, 5)
    beta1_prior_mean = 1
    beta1_prior_var = 5

    # σ² ~ Inverse-Gamma(2, 1)
    sigma2_prior_a = 2
    sigma2_prior_b = 1

    # 3. 定义重要抽样分布
    # 使用最小二乘进行估计
    X = np.column_stack([np.ones(n), x])
    beta_estimate = np.linalg.inv(X.T @ X) @ X.T @ y
    y_pred = X @ beta_estimate
    residuals = y - y_pred
    RSS = residuals.T @ residuals
    sigma2_estimate = RSS / (n - 2)

    print(f"最小二乘结果：β0 = {beta_estimate[0]}, β1 = {beta_estimate[1]}, σ² = {sigma2_estimate}")
    print(f"残差平方和 RSS = {RSS}")

    

    # 4. 定义联合后验分布(未归一化)
    def unnormalized_log_posterior(param):
        """
        计算后验概率密度的未归一化值

        param: [β0, β1, log_sigma2)]
        """

        beta0, beta1, log_sigma2 = param

        sigma2 = np.exp(log_sigma2)

        # 计算似然函数值
        y_pred = beta0 + beta1 * x
        residuals = y - y_pred
        log_likelihood = -0.5 * n * np.log(2 * np.pi * sigma2) * -0.5 * np.sum(residuals**2) / sigma2

        # 计算先验概率
        log_prior_beta0 = stats.norm.logpdf(beta0, loc = beta0_prior_mean, scale = np.sqrt(beta0_prior_var))
        log_prior_beta1 = stats.norm.logpdf(beta1, loc = beta1_prior_mean, scale = np.sqrt(beta1_prior_var))
        log_prior_sigma2 = stats.invgamma.logpdf(sigma2, a = sigma2_prior_a, scale = sigma2_prior_b) + log_sigma2

        return log_likelihood + log_prior_beta0 + log_prior_beta1 + log_prior_sigma2
    
    # 5. 定义重要抽样分布
    proposal_mean = np.array([beta_estimate[0], beta_estimate[1], np.log(sigma2_estimate)])
    proposal_cov = np.diag([1.0, 1.0, 1.0])

    def log_proposal_density(param):
        return stats.multivariable_normal.logpdf(param, mean = proposal_mean, cov = proposal_cov)
    
    # 6. 进行重要抽样
    num_samples = 10000
    samples = np.zeros((num_samples, 3))
    weights = np.zeros(num_samples)
    log_weights = np.zeros(num_samples)

    for i in range(num_samples):
        samples[i] = np.random.multivariate_normal(proposal_mean, proposal_cov)
        unnormalized_log_posterior = unnormalized_log_posterior(samples[i])
        log_proposal = log_proposal_density(samples[i])
        log_weights[i] = unnormalized_log_posterior - log_proposal
 
    weights = np.exp(log_weights)



if __name__ == "__main__":
    baysian_linear_regression_inportance_sampling()