Baysian_Canonical_Identification/RJMCMC_ModelSelector.py

import numpy as np
from scipy import stats


class RJMCMC_Sampler:
    def __init__(self, y, u, k_min=1, k_max=10, initial_k=None,
                 initial_a=None, initial_b=None):
        """
        初始化采样器

        参数:
        y (array): 观测到的输出序列
        u (array): 已知的输入序列
        k_min (int): 探索的最小模型阶数
        k_max (int): 探索的最大模型阶数
        initial_k (int, optional): 初始模型阶数
        initial_a (dict, optional): 初始多项式系数
        initial_b (array, optional): 初始多项式系数
        """

        # 定义观测数据
        self.y = y
        self.u = u
        self.N = len(y)

        # 定义模型阶数
        self.k_min = k_min
        self.k_max = k_max
        gamma_sample = np.random.gamma(shape=2, scale=1)
        self.current_k = initial_k if initial_k is not None else min(max(int(gamma_sample), k_min), k_max)

        # 初始化多项式系数
        if initial_a is not None:
            self.current_a = initial_a
        else:
            self.current_a = self._generate_stable_a(self.current_k)

        # 定义标准型的b系数
        if initial_b is not None:
            self.current_b = initial_b
        else:
            self.current_b = self._generate_observable_b(self.current_k)

        # 初始化噪声先验 (假设为 Gamma 先验)
        self.lambda_a_w = 1e-3 # 过程噪声精度的先验
        self.lambda_b_w = 1e-3
        self.lambda_a_z = 1e-3 # 观测噪声精度的先验
        self.lambda_b_z = 1e-3

        # 初始化噪声方差 (从先验采样)
        # 为避免除零错误，在分母上增加一个极小值
        epsilon = 1e-9
        self.current_sigma2_w = 1.0 / (np.random.gamma(self.lambda_a_w, 1.0/self.lambda_b_w) + epsilon)
        self.current_sigma2_z = 1.0 / (np.random.gamma(self.lambda_a_z, 1.0/self.lambda_b_z) + epsilon)


    def _update_current_eigenvalues(self):
        """一个辅助函数，根据 current_a 更新 current_eigenvalues。"""
        if self.current_k > 0:
            poly_coeffs = np.concatenate(([1], -self.current_a[::-1]))
            self.current_eigenvalues = np.roots(poly_coeffs)
        else:
            self.current_eigenvalues = np.array([])

    def _update_current_a_coeffs(self):
        """一个辅助函数，根据 current_eigenvalues 更新 current_a。"""
        k = len(self.current_eigenvalues)
        if k > 0:
            # np.poly 返回 [1, a_{k-1}, ..., a_0]，去掉首项“1”，反转剩下的
            poly_coeffs = np.poly(self.current_eigenvalues)
            self.current_a = np.real(poly_coeffs[1:][::-1])
        else:
            self.current_a = np.array([])
    
    def _cal_current_eigenvalues(self, a_coeffs_temp):
        """一个辅助函数，根据给定的 a_coeffs 计算对应的特征值。"""
        k = len(a_coeffs_temp)
        if k > 0:
            poly_coeffs = np.concatenate(([1], -a_coeffs_temp[::-1]))
            return np.roots(poly_coeffs)
        else:
            return np.array([])
        
    def _cal_current_a_coeffs(self, eigenvalues_temp):
        """一个辅助函数，根据给定的特征值计算对应的 a_coeffs。"""
        k = len(eigenvalues_temp)
        if k > 0:
            # np.poly 返回 [1, a_{k-1}, ..., a_0]，去掉首项“1”，反转剩下的
            poly_coeffs = np.poly(eigenvalues_temp)
            return np.real(poly_coeffs[1:][::-1])
        else:
            return np.array([])


    def _log_prior_eigenvalues(self, eigenvalues):
        """
        计算 k 个特征值 (Lambda) 的对数先验概率 log p(Lambda)。
        
        该先验基于论文 4.1.1 节 讨论的先验：
        1. 稳定性：如果任何 |lambda| >= 1，先验概率为 0 (log(0) = -inf)。
        2. 实数根 (p_real)：在 (-1, 1) 上均匀分布, p(lambda_r) = 1/2。
           log p(lambda_r) = -log(2)。
        3. 复数根 (p_complex)：在单位圆盘上均匀分布, p(lambda_c) = 1/pi。
           log p(lambda_c) = -log(pi)。
           
        注意：复共轭对 (lambda, lambda_conj) 只计算一次 (例如上半平面)。
              我们通过只计算 imag(lambda) >= 0 的复数根来实现。
        """
        
        log_prior = 0.0
        
        # 跟踪已处理的复数根 (避免重复计算共轭对)
        processed_complex = set()

        for eig in eigenvalues:
            if np.abs(eig) >= 1.0:
                # 违反稳定性约束
                return -np.inf
            
            if np.isclose(eig.imag, 0):
            # 是实数
            # 这是 "Uniform real-eigenvalue prior"
                log_prior += -np.log(2.0)

            else:
                # 是复数
                # 检查是否已处理过
                if eig in processed_complex or np.conjugate(eig) in processed_complex:
                    continue
                
                # 我们只计算 imag >= 0 的根 (代表一对)
                if eig.imag >= 0:
                    # 这是 "Polar-coordinate prior"
                    log_prior += -np.log(np.pi) 
                    processed_complex.add(eig)
                
        return log_prior
    
    def _log_prior_eigenvalues_to_a(self,eigenvalues):
        """
        计算从特征值 (eigenvalues) 到多项式系数 a = [a_0, ..., a_{k-1}] 的对数先验概率。
        
        该方法基于 Proposition 4.2 (Change of variables) 
        log p(a) = log p(Lambda) - log(|det J|)
        
        其中 log p(Lambda) 是特征值的先验
        而 log(|det J|) 是维塔变换的对数雅可比行列式。
        """
        # 计算 log p(Lambda)
        log_prior_lambda = self._log_prior_eigenvalues(eigenvalues)

        # 如果特征值不稳定 (log p(Lambda) = -inf)，则 p(a) 也不可能
        if log_prior_lambda == -np.inf:
            return -np.inf
        
        # 计算对数雅可比行列式 log(|det J|)
        log_det_jacobian = 0.0
        k = len(eigenvalues)
        if k < 2:
            return log_prior_lambda
        
        for i in range(k):
            
            for j in range(i + 1, k):
                diff = eigenvalues[i] - eigenvalues[j]
                abs_diff = np.abs(diff)
                if abs_diff == 0:
                    return -np.inf  # 重根导致雅可比行列式为零
                log_det_jacobian += np.log(abs_diff)
            
            if log_det_jacobian == -np.inf:
                return -np.inf
            
        # 应用变量替换公式
        log_prior_a = log_prior_lambda - log_det_jacobian

        return log_prior_a
    
    def _log_likelihood(self, k, a_coeffs, b_coeffs, sigma2_w, sigma2_z):
        """
        计算给定参数下的对数似然 log p(y | k, a, b, sigma_w, sigma_z)。
        
        使用标准卡尔曼滤波器 (Kalman Filter) (如论文 Appendix B ) 
        和预测误差分解 (Prediction Error Decomposition) (如论文 Eq. (D.37) )。
        """
        
        # 如果 k=0，模型无法运行
        if k == 0:
            # 返回一个非常小的似然
            return -np.inf
        
        # 获取当前参数下的状态空间矩阵（以参数为输入，避免依赖实例状态）
        A, B, C, D = self.get_controller_canonical_form(k, a_coeffs, b_coeffs)
        
        # 过程噪声 Sigma (k x k)
        Sigma = np.zeros((k, k))
        Sigma[k-1, k-1] = sigma2_w

        # 测量噪声 Gamma (1 x 1)
        Gamma = np.array([[sigma2_z]])

        # 初始化卡尔曼滤波器
        
        # 状态 x_t|t-1
        x_pred = np.zeros((k, 1))

        # 状态协方差 P_t|t-1
        P_pred = np.eye(k) * 1e6

        total_log_likelihood = 0.0

        # 运行卡尔曼滤波器 N 步
        for t in range(self.N):
            y_t = self.y[t]
            u_t = self.u[t]
            # --- 预测步 (Prediction) ---
            # (在 t=0 时，x_pred 和 P_pred 是我们初始化的 x_0|-1, P_0|-1)

            # --- 计算似然 ---
            # 预测误差 nu_t 
            y_pred = C @ x_pred + D @ np.array([[u_t]])
            nu_t = y_t - y_pred

            # 预测误差协方差 S_t
            S_t = C @ P_pred @ C.T + Gamma
            S_t_inv = 1.0 / S_t

            # 累加对数似然
            total_log_likelihood += -0.5 * (np.log(2 * np.pi) + np.log(S_t) + (nu_t * S_t_inv * nu_t))

            # --- 更新步 (Update) ---
            # 卡尔曼增益 K_t (k x 1)
            K_t = P_pred @ C.T * S_t_inv

            # 更新状态 x_t|t
            x_update = x_pred + K_t @ nu_t

            # 更新协方差 P_t|t
            I_KC = np.eye(k) - K_t @ C
            P_update = I_KC @ P_pred @ I_KC.T + K_t @ Gamma @ K_t.T

            # --- 为下一次循环准备预测 (t+1) ---
            x_pred = A @ x_update + B @ np.array([[u_t]])
            P_pred = A @ P_update @ A.T + Sigma
        
        return total_log_likelihood # 返回标量值
    
    def _propose_birth(self):
        """
        提议一个 "诞生" 转移 (k -> k+1 或 k -> k+2)。
        返回包含提议状态和计算接受率所需的对数比率的字典。
        """
        k = self.current_k
        current_eigs = self.current_eigenvalues

        # 获取实数特征值数目
        real_eigs_indices = np.where(np.isreal(current_eigs))
        num_real_eigs = len(real_eigs_indices[0])

        # 获取复数特征值数目，只保留 imag > 0 的部分 (共轭对只算一次)
        complex_eigs_indices = np.where((np.iscomplex(current_eigs)) & (current_eigs.imag > 0)) #type: ignore
        num_complex_eigs = len(complex_eigs_indices[0])

        # 决定是诞生一个实数根还是一对复共轭根
        can_add_complex = (k + 2) <= self.k_max
        add_real = True
        if can_add_complex:
        # 以0.5的概率选择诞生一对复共轭根
            if np.random.rand() < 0.5:
                add_real = False

        if add_real:
            k_new = k + 1

            # 诞生实根的对数概率之比
            log_birth_real_forward = np.log(0.5)
            log_birth_real_backward = np.log(0.5)

            # 从提议分布 q(u) 中采样辅助变量 u = (u_lambda, u_b)
            u_lambda = np.random.uniform(-1.0, 1.0)  # 新特征值
            u_b = np.random.normal(0, 1)             # 新 b 系数

            # 计算对数提议密度
            log_q_forward = log_birth_real_forward + stats.norm.logpdf(u_b, 0, 1) - np.log(2.0)
            log_q_backward = log_birth_real_backward - np.log(num_real_eigs + 1)

            # 计算对数雅可比行列式 log|J|
            if k == 0:
                log_det_jacobian = 0.0
            else:
                log_det_jacobian = 1

            # 构造新状态
            new_eigs = np.append(current_eigs, u_lambda)
            new_a = np.real(np.poly(new_eigs)[1:][::-1])
            new_b = np.append(self.current_b, u_b)

            # 接受率中的项为 q_backward(u) * |J| / q_forward(u)
            log_ratio = log_q_backward + log_det_jacobian - log_q_forward

            # 计算未归一化的后验之比
            log_post_unnormalized_current = self._log_posterior_unnormalized(k, self.current_a, self.current_b,
                                                                          self.current_sigma2_w, self.current_sigma2_z)
            log_post_unnormalized_new = self._log_posterior_unnormalized(k_new, new_a, new_b, self.current_sigma2_w, self.current_sigma2_z)

            # 计算接受率
            log_acceptance_ratio = log_post_unnormalized_new - log_post_unnormalized_current + log_ratio
            acceptance_ratio = np.exp(log_acceptance_ratio)

            return {"k_new": k_new, "a_new": new_a, "b_new": new_b, "log_ratio": log_ratio, "type": "birth_real", "new_eigs": new_eigs, "acceptance_ratio": acceptance_ratio}

        else:
            # --- 诞生一对复共轭根 (k -> k+2) ---
            k_new = k + 2

            # 诞生实根的对数概率之比
            log_birth_complex_forward = np.log(0.5)
            log_birth_complex_backward = np.log(0.5)

            # 采样辅助变量 u = (rho, theta, u_b1, u_b2)
            rho = np.sqrt(np.random.uniform(0, 1.0))
            theta = np.random.uniform(0, np.pi)
            u_lambda = rho * (np.cos(theta) + 1j * np.sin(theta))
            u_b1, u_b2 = np.random.normal(0, 1, 2)

            # 计算对数提议密度 log(q(u))
            log_q_forward = log_birth_complex_forward + stats.norm.logpdf(u_b1, 0, 1) + stats.norm.logpdf(u_b2, 0, 1) - np.log(np.pi) - np.log(1.0)
            log_q_backward = log_birth_complex_backward - np.log(num_complex_eigs + 1)

            # 计算对数雅可比行列式 log|J|
            # |J| = 2*ρ
            log_jacobian = np.log(2 * rho)
            
            # 构造新状态
            new_eigs = np.append(current_eigs, [u_lambda, np.conjugate(u_lambda)])
            new_a = np.real(np.poly(new_eigs)[1:][::-1])
            new_b = np.append(self.current_b, [u_b1, u_b2])

            log_ratio = log_q_backward + log_jacobian - log_q_forward

            # 计算未归一化的后验之比
            log_post_unnormalized_current = self._log_posterior_unnormalized(k, self.current_a, self.current_b,
                                                                          self.current_sigma2_w, self.current_sigma2_z)
            log_post_unnormalized_new = self._log_posterior_unnormalized(k_new, new_a, new_b, self.current_sigma2_w, self.current_sigma2_z)
            
            # 计算接受率
            log_acceptance_ratio = log_post_unnormalized_new - log_post_unnormalized_current + log_ratio
            acceptance_ratio = np.exp(log_acceptance_ratio)

            return {"k_new": k_new, "a_new": new_a, "b_new": new_b, "log_ratio": log_ratio, "type": "birth_complex", "new_eigs": new_eigs, "acceptance_ratio": acceptance_ratio}

    def _propose_death(self):
        """
        提议一个 "消亡" 转移 (k -> k-1 或 k -> k-2)。
        返回包含提议状态和计算接受率所需的对数比率的字典。
        """
        k = self.current_k
        current_eigs = self.current_eigenvalues

        # 识别出可移除的实数根和复共轭对
        real_eigs_indices = np.where(np.isreal(current_eigs))
        complex_eigs_indices = np.where((np.iscomplex(current_eigs)) & (current_eigs.imag > 0)) #type: ignore
    
        can_remove_real = len(real_eigs_indices[0]) > 0
        can_remove_complex = len(complex_eigs_indices[0]) > 0

        if not can_remove_real and not can_remove_complex:
            return None # 无法执行消亡
        
        # 决定是移除实数根还是复共轭对
        remove_real = True
        if can_remove_real and can_remove_complex:
            if np.random.rand() < 0.5:
                remove_real = False
        elif can_remove_complex:
            remove_real = False

        if remove_real:
         # --- 消亡一个实数根 (k -> k-1) ---
            k_new = k - 1

            # 生成与消亡实根对应的诞生过程的对数概率之比
            log_death_real_forward = np.log(0.5)
            log_death_real_backward = np.log(0.5)

            # 1. 随机选择一个实数根移除
            idx_to_remove = np.random.choice(real_eigs_indices[0])
            lambda_removed = current_eigs[idx_to_remove]

            # 2. 随机选择一个实数 b 系数移除
            idx_b_to_remove = np.random.choice(len(self.current_b))
            b_removed = self.current_b[idx_b_to_remove]

            # 移除的根和b系数构成了逆向（诞生）提议的辅助变量 u
            u_lambda = lambda_removed
            u_b = b_removed

            # 2. 计算逆向提议的对数密度 log(q(u))
            log_q_forward = log_death_real_forward - np.log(len(real_eigs_indices[0])) - np.log(len(self.current_b))
            log_q_backward = log_death_real_backward + stats.norm.logpdf(u_b, 0, 1) - np.log(2.0)
 
            # 3. 计算对应诞生过程的对数雅可比行列式 log|J|
            remaining_eigs = np.delete(current_eigs, idx_to_remove)
            if k_new == 0:
                log_jacobian_birth = 0.0
            else:
                log_jacobian_birth = 1

            # 构造新状态
            if k_new == 0:
                new_a = np.array([])
            else:
                new_a = np.real(np.poly(remaining_eigs)[1:][::-1])
            new_b = np.delete(self.current_b, idx_b_to_remove)

            # 接受率中的项为 q_backward(u) / (|J| * q_forward(u))
            log_ratio = -log_q_forward - log_jacobian_birth + log_q_backward

            # 计算未归一化的后验之比
            log_post_unnormalized_current = self._log_posterior_unnormalized(k, current_eigs, self.current_b, self.current_sigma2_w, self.current_sigma2_z)
            log_post_unnormalized_new = self._log_posterior_unnormalized(k_new, remaining_eigs, new_b, self.current_sigma2_w, self.current_sigma2_z)

            # 计算接受率
            log_acceptance_ratio = log_post_unnormalized_new - log_post_unnormalized_current + log_ratio
            acceptance_ratio = np.exp(log_acceptance_ratio)

            return {"k_new": k_new, "a_new": new_a, "b_new": new_b, "log_ratio": log_ratio, "type": "death_real", "new_eigs": remaining_eigs, "acceptance_ratio": acceptance_ratio}

        else:
            # --- 消亡一对复共轭根 (k -> k-2) ---
            k_new = k - 2
            rho = np.sqrt(np.random.uniform(0, 1.0))

            # 生成与消亡复共轭对对应的诞生过程的对数概率之比
            log_death_complex_forward = np.log(0.5)
            log_death_complex_backward = np.log(0.5)
            
            # 随机选择一对共轭根移除
            complex_idx_to_remove = np.random.choice(complex_eigs_indices[0])
            lambda_removed = current_eigs[complex_idx_to_remove]

            # 随机选择一对 b 系数移除，需要确保两次选的不一样
            idx_b1_to_remove, idx_b2_to_remove = np.random.choice(len(self.current_b), size=2, replace=False)
            b1_removed = self.current_b[idx_b1_to_remove]
            b2_removed = self.current_b[idx_b2_to_remove]

            # 找到其共轭对
            conjugate_idx_to_remove_array = np.where(current_eigs == np.conjugate(lambda_removed))[0]
            if len(conjugate_idx_to_remove_array) == 0:
                return None  # 找不到共轭对，无法执行消亡
            conjugate_idx = conjugate_idx_to_remove_array[0]
            
            # 逆向提议的辅助变量 u
            u_lambda = lambda_removed
            u_b1, u_b2 = self.current_b[-2:]
            
            # 计算逆向提议的对数密度 log(q(u))
            n_b = len(self.current_b)
            pair_count = n_b * (n_b - 1) / 2.0
            log_q_forward = log_death_complex_forward - np.log(len(complex_eigs_indices[0])) - np.log(pair_count)
            log_q_backward = log_death_complex_backward + stats.norm.logpdf(u_b1, 0, 1) + stats.norm.logpdf(u_b2, 0, 1) - np.log(np.pi) - np.log(1.0)
            
            # 计算对应诞生过程的对数雅可比行列式
            remaining_eigs = np.delete(current_eigs, [complex_idx_to_remove, conjugate_idx])
            log_jacobian = np.log(2 * rho)
            
            # 构造新状态
            if k_new == 0:
                new_a = np.array([])
            else:
                new_a = np.real(np.poly(remaining_eigs)[1:][::-1])
            new_b = np.delete(self.current_b, [idx_b1_to_remove, idx_b2_to_remove])
            
            log_ratio = -log_q_forward - log_jacobian + log_q_backward

            # 计算未归一化的后验之比
            log_post_unnormalized_current = self._log_posterior_unnormalized(k, current_eigs, self.current_b, self.current_sigma2_w, self.current_sigma2_z)
            log_post_unnormalized_new = self._log_posterior_unnormalized(k_new, remaining_eigs, new_b, self.current_sigma2_w, self.current_sigma2_z)

            # 计算接受率
            log_acceptance_ratio = log_post_unnormalized_new - log_post_unnormalized_current + log_ratio
            acceptance_ratio = np.exp(log_acceptance_ratio)

            return {"k_new": k_new, "a_new": new_a, "b_new": new_b, "log_ratio": log_ratio, "type": "death_complex", "new_eigs": remaining_eigs, "acceptance_ratio": acceptance_ratio}

    def _propose_within_model(self):
        """
        随机选择实根变复根，复根变实根，或微调现有根，调用现有函数_merge_eigenvalues，_split_eigenvalues，_disturbe_swimming。
        """
        k = self.current_k
        current_eigs = self.current_eigenvalues.copy()

        if k == 0:
            return None  # 无法在 k=0 时进行模型内提议
        
        proposal_type = np.random.choice(['merge', 'split', 'disturb'], p=[0.3, 0.3, 0.4])

        if proposal_type == 'merge':
            return self._merge_eigenvalues(current_eigs)  # type: ignore
        elif proposal_type == 'split':
            return self._split_eigenvalues(current_eigs)  # type: ignore
        else:
            return self._disturbe_swimming(current_eigs)  # type: ignore
        

    def _log_prior_b(self, b_coeffs):
        """
        计算 C_c 系数 (b_0, ..., b_{k-1}) 的对数先验概率。
        根据论文 4.2 节，我们使用独立标准正态先验。
        b_i ~ N(0, 1)
        """
        #
        return stats.norm.logpdf(b_coeffs, 0, 1).sum()
    

    def _generate_stable_a(self, k):
            """
            生成 k 个稳定的 a = [a_0, a_1, ..., a_{k-1}] 系数, 并保存对应的特征值。
            
            该方法基于论文 4.1.1 节 [cite: 314-320] 和 6.2 节  的讨论，
            生成 k 个稳定的特征值 (根)，然后使用维塔定理 (Proposition 4.1) [cite: 290] 
            计算多项式系数。
            
            为确保系数a为实数，特征值必须是实数或复共轭对 。
            """
            
            eigenvalues = []
            i = 0
            
            while i < k:
                # 如果 k-i > 1，随机决定采样一个实数根还是一个复共轭对
                # 如果 k-i == 1，则只能采样一个实数根
                if i == k - 1 or np.random.rand() < 0.5:
                    # --- 1. 采样一个实数根 ---
                    # 使用 "Uniform real-eigenvalue prior" (在 (-1, 1) 上均匀分布) 
                    real_eig = np.random.uniform(-1.0, 1.0)
                    eigenvalues.append(real_eig)
                    i += 1
                else:
                    # --- 2. 采样一个复共轭对 ---
                    # 使用 "Polar-coordinate prior" (极坐标先验) 
                    # 为确保在单位圆盘内面积均匀，rho = sqrt(U(0,1)) 
                    rho = np.sqrt(np.random.uniform(0, 1.0))
                    
                    # theta 在 (0, pi) 均匀分布 (上半平面) 
                    theta = np.random.uniform(0, np.pi)
                    
                    # 计算复特征值
                    complex_eig = rho * (np.cos(theta) + 1j * np.sin(theta))
                    
                    # 添加该值及其共轭
                    eigenvalues.append(complex_eig)
                    eigenvalues.append(np.conjugate(complex_eig))
                    i += 2
            
            # --- 3. (维塔定理) 从根计算多项式系数 ---
            # np.poly(roots) 返回 [1, a_{k-1}, a_{k-2}, ..., a_0]
            coefficients_descending = np.poly(eigenvalues)
          
            # 返回 [a_0, a_1, ..., a_{k-1}]
            a_coefficients = coefficients_descending[1:][::-1]
            
            self.current_eigenvalues = np.array(eigenvalues) # 顺便保存特征值
            return np.real(a_coefficients)


    def _generate_observable_b(self, k):
        """
        生成 b = [b_0, ..., b_{k-1}] 系数 (共 k 个)。
        """
        b_array = np.random.normal(0, 1, size=k)
        return b_array
    

    def get_controller_canonical_form(self, k=None, a_coeffs=None, b_coeffs=None):
        """
        返回给定参数的规范型 (Controller Canonical Form) 矩阵 A_c, B_c, C_c, D_c。

        如果提供了 k、a_coeffs、b_coeffs，则使用这些输入（使函数成为纯函数）；
        否则回退到实例的当前状态（向后兼容）。
        """
        # 回退到实例状态以保持向后兼容
        if k is None:
            k = self.current_k
        if a_coeffs is None:
            a_coeffs = self.current_a
        if b_coeffs is None:
            b_coeffs = self.current_b

        # k=0 是无效情况，但 k=1 时 np.diag(np.ones(0), 1) 会创建一个空 1x1 矩阵
        if k == 0:
            return np.array([]), np.array([]), np.array([]), np.array([[0.0]])

        # 1. 构造 A_c (k x k)
        # 先创建一个 k x k 的零矩阵
        A_c = np.zeros((k, k))
        if k > 1:
            # 创建上对角线为 1
            np.fill_diagonal(A_c[0:-1, 1:], 1)

        # 填充最后一行
        # 确保 a_coeffs 的形状/长度与 k 匹配
        a_arr = np.asarray(a_coeffs).ravel()
        if a_arr.size != k:
            raise ValueError(f"a_coeffs length {a_arr.size} does not match k={k}")
        A_c[-1, :] = -a_arr

        # 2. 构造 B_c (k x 1)
        B_c = np.zeros((k, 1))
        B_c[-1] = 1

        # 3. 构造 C_c (1 x k)
        b_arr = np.asarray(b_coeffs).ravel()
        if b_arr.size != k:
            raise ValueError(f"b_coeffs length {b_arr.size} does not match k={k}")
        C_c = b_arr.reshape(1, k)

        # 4. 构造 D_c (1 x 1)
        # 论文 Definition 3.1 包含 d_0 ，但实验中设为 0 
        D_c = np.array([[0.0]])

        return A_c, B_c, C_c, D_c
    
    def _log_posterior_unnormalized(self, k, eigenvalues, b_coeffs, sigma2_w, sigma2_z):
        """计算所有参数的非归一化对数后验。"""
        log_prior_k = -np.log(self.k_max - self.k_min + 1)
        
        if k > 0:
            log_prior_eigen = self._log_prior_eigenvalues(eigenvalues)
        else:
            log_prior_eigen = 0.0

        log_prior_b = self._log_prior_b(b_coeffs)
        precision_w = 1.0 / sigma2_w
        precision_z = 1.0 / sigma2_z
        log_prior_w = stats.gamma.logpdf(precision_w, a=self.lambda_a_w, scale=1.0/self.lambda_b_w)
        log_prior_z = stats.gamma.logpdf(precision_z, a=self.lambda_a_z, scale=1.0/self.lambda_b_z)
        log_likelihood = self._log_likelihood(k, self._cal_current_a_coeffs(eigenvalues), b_coeffs, sigma2_w, sigma2_z)
        return log_prior_k + log_prior_eigen + log_prior_b + log_prior_w + log_prior_z + log_likelihood
    
    def _purpose_b_swimming(self):
        """
        对 b 系数进行游动：对每个 b_i 添加一个小的高斯扰动, 返回b游动前后的先验
        """
        k = self.current_k
        current_b = self.current_b
        proposal_b = np.copy(current_b)

        # 对每个 b_i 添加高斯扰动
        for i in range(k):
            b_proposal = proposal_b[i] + np.random.normal(0, 0.1)  # 小的高斯扰动
            proposal_b[i] = b_proposal

        # 计算新b的先验
        log_prior_current = self._log_prior_b(current_b)
        log_prior_proposal = self._log_prior_b(proposal_b)

        return proposal_b, log_prior_current, log_prior_proposal
    
    def _purpose_lumbda_swimming(self):
        """
        对 lumbda 进行游动：对每个 lumbda 添加一个小的高斯扰动，虚部也要有扰动
        """
        k = self.current_k
        current_eigs = np.copy(self.current_eigenvalues)
        proposal_eigs = np.copy(current_eigs)

        # 对每个 lumbda 添加高斯扰动
        for i in range(k):
            eig_proposal = proposal_eigs[i] + np.random.normal(0, 0.1)  # 小的高斯扰动
            proposal_eigs[i] = eig_proposal

        # 虚部也要有扰动，共轭的虚部记得相等
        for i in range(k):
            if np.iscomplex(proposal_eigs[i]) and proposal_eigs[i].imag != 0:
                imag_perturbation = np.random.normal(0, 0.1)
                proposal_eigs[i] = proposal_eigs[i].real + 1j * (proposal_eigs[i].imag + imag_perturbation)
                # 找到共轭根并更新
                conjugate_idx = np.where(proposal_eigs == np.conjugate(current_eigs[i]))[0]
                if len(conjugate_idx) > 0:
                    proposal_eigs[conjugate_idx[0]] = proposal_eigs[conjugate_idx[0]].real - 1j * (current_eigs[i].imag + imag_perturbation)

        # 计算新lumbda的先验
        log_prior_current = self._log_prior_eigenvalues(current_eigs)
        log_prior_proposal = self._log_prior_eigenvalues(proposal_eigs)

        return proposal_eigs, log_prior_current, log_prior_proposal

    
    
    def _merge_eigenvalues(self, current_eigs):
        """合并游动：选择两个实数根，确定性地合并为一个复共轭对。"""
        k = self.current_k
    
        # 随机选择两个不同的实数根(可能是虚部为0的复数)，并删除他们
        real_indices = np.where(np.isreal(current_eigs))
        if len(real_indices[0]) < 2:
            return  # 不足两个实数根，无法合并
        idx_pair = np.random.choice(real_indices[0], size=2, replace=False)
        idx1, idx2 = idx_pair
        lambda1 = current_eigs[idx1].real   # type: ignore
        lambda2 = current_eigs[idx2].real   # type: ignore

        current_eigs = np.delete(current_eigs, [idx1, idx2])

        # 生成一对共轭复根
        new_theta = np.random.uniform(0, np.pi)
        new_rho = np.random.uniform(0, 1.0)
        new_eigenvalue = new_rho * (np.cos(new_theta) + 1j * np.sin(new_theta))
        current_eigs = np.append(current_eigs, [new_eigenvalue, np.conjugate(new_eigenvalue)])

        # 对b进行游动
        proposal_b, log_prior_b_current, log_prior_b_proposal = self._purpose_b_swimming()

        # 计算提议密度
        # 正向过程是任选两个实根，合并为复共轭对
        log_q_merge_forward = -np.log(len(real_indices[0]) * (len(real_indices[0]) - 1) / 2.0) - np.log(np.pi) - np.log(1.0)  # 提议密度 q_forward(u) 在 (0, pi) x (0,1)
        # 逆向过程是从复共轭对中任选一个删去，然后抽样两个新实根,提议分布就是均匀分布
        num_complex_eigs = len(np.where((np.iscomplex(current_eigs)) & (current_eigs.imag > 0))[0])  #type: ignore
        log_q_merge_backward = -np.log(num_complex_eigs + 1) - np.log(2) - np.log(2)

        # 计算对数提议比
        log_proposal_ratio = log_q_merge_backward - log_q_merge_forward

        # 计算后验比
        log_post_unnormalized_current = self._log_posterior_unnormalized(k, self.current_eigenvalues, self.current_b, self.current_sigma2_w, self.current_sigma2_z)
        log_post_unnormalized_proposal = self._log_posterior_unnormalized(k, current_eigs, proposal_b, self.current_sigma2_w, self.current_sigma2_z)

        # 计算接受率
        log_acceptance_ratio = (log_post_unnormalized_proposal - log_post_unnormalized_current) + log_proposal_ratio

        acceptance_ratio = np.exp(log_acceptance_ratio)

        return {"a_new": np.real(np.poly(current_eigs)[1:][::-1]), "b_new": proposal_b, "log_ratio": log_proposal_ratio, "type": "merge", "new_eigs": current_eigs, "acceptance_ratio": acceptance_ratio}
    

    def _split_eigenvalues(self, current_eigs):
        """
        分裂游动：选择一个复共轭对，随机地分裂为两个实数根。 
        """
        k = self.current_k

        # 随机选择一个复共轭对，并删除它们
        complex_indices = np.where((np.iscomplex(current_eigs)) & (current_eigs.imag > 0)) #type: ignore
        if len(complex_indices[0]) < 1:
            return  # 没有复共轭对，无法分裂
        idx_to_remove = np.random.choice(complex_indices[0])
        lambda_removed = current_eigs[idx_to_remove]
        conjugate_idx_to_remove = np.where(current_eigs == np.conjugate(lambda_removed))[0]
        
        if len(conjugate_idx_to_remove) == 0:
            return  # 找不到共轭对，无法分裂
        conjugate_idx = conjugate_idx_to_remove[0]

        # 删除选中的复共轭对
        current_eigs = np.delete(current_eigs, [idx_to_remove, conjugate_idx])
        # 生成两个实数根
        real_eig1 = np.random.uniform(-1.0, 1.0)
        real_eig2 = np.random.uniform(-1.0, 1.0)
        current_eigs = np.append(current_eigs, [real_eig1, real_eig2])

        # 对b进行游动
        proposal_b, log_prior_b_current, log_prior_b_proposal = self._purpose_b_swimming()

        # 计算提议密度
        # 正向过程是任选一个复共轭对，分裂为两个实根
        log_q_split_forward = -np.log(len(complex_indices[0])) - np.log(2) - np.log(2)  # 提议密度 q_forward(u) 在 (-1,1) x (-1,1)
        # 逆向过程是从实根中任选两个删去，然后抽样一个复共轭对,提议分布就是均匀分布
        num_real_eigs = len(np.where(np.isreal(current_eigs))[0])
        log_q_split_backward = -np.log(num_real_eigs * (num_real_eigs - 1) / 2.0) - np.log(np.pi) - np.log(1.0)

        # 计算对数提议比
        log_proposal_ratio = log_q_split_backward - log_q_split_forward

        # 计算后验比
        log_post_unnormalized_current = self._log_posterior_unnormalized(k, self.current_eigenvalues, self.current_b, self.current_sigma2_w, self.current_sigma2_z)
        log_post_unnormalized_proposal = self._log_posterior_unnormalized(k, current_eigs, proposal_b, self.current_sigma2_w, self.current_sigma2_z)

        # 计算接受率
        log_acceptance_ratio = (log_post_unnormalized_proposal - log_post_unnormalized_current) + log_proposal_ratio
        acceptance_ratio = np.exp(log_acceptance_ratio)

        return {"a_new": np.real(np.poly(current_eigs)[1:][::-1]), "b_new": proposal_b, "log_ratio": log_proposal_ratio, "type": "split", "new_eigs": current_eigs, "acceptance_ratio": acceptance_ratio}
    
    def _disturbe_swimming(self, current_eigs):
        """
        扰动游动：保持维度不变，保持实数根和复共轭对的数量不变，进行扰动。
        """
        k = self.current_k

        # 对特征值进行游动
        proposal_eigs, log_prior_eigs_current, log_prior_eigs_proposal = self._purpose_lumbda_swimming()

        # 对b进行游动
        proposal_b, log_prior_b_current, log_prior_b_proposal = self._purpose_b_swimming()

        # 计算提议密度比 (对称提议，密度相等)
        log_proposal_ratio = 0.0

        # 计算后验比
        log_post_unnormalized_current = self._log_posterior_unnormalized(k, self.current_eigenvalues, self.current_b, self.current_sigma2_w, self.current_sigma2_z)
        log_post_unnormalized_proposal = self._log_posterior_unnormalized(k, proposal_eigs, proposal_b, self.current_sigma2_w, self.current_sigma2_z)

        # 计算接受率
        log_acceptance_ratio = (log_post_unnormalized_proposal - log_post_unnormalized_current) + log_proposal_ratio
        acceptance_ratio = np.exp(log_acceptance_ratio)

        return {"a_new": np.real(np.poly(proposal_eigs)[1:][::-1]), "b_new": proposal_b, "log_ratio": log_proposal_ratio, "type": "disturbe", "new_eigs": proposal_eigs, "acceptance_ratio": acceptance_ratio}
    

    def run_MCMC(self, num_iterations):
        """
        运行 RJMCMC 采样器指定次数的迭代。
        """
        # 存储采样历史
        k_history = []
        a_history = []
        b_history = []
        sigma2_w_history = []
        sigma2_z_history = []
        eigenvalues_history = []

        print(f"Starting MCMC with initial k={self.current_k}")

        for i in range(num_iterations):
            # 随机选择一种转移类型
            # 这里的概率可以根据需要调整
            move_type = np.random.choice(['birth_death', 'within_model'], p=[0.5, 0.5])

            proposal = None
            accepted = False
            move_name = 'None'

            if move_type == 'birth_death':
                # 决定是 birth 还是 death
                if self.current_k == self.k_min:
                    proposal = self._propose_birth()
                elif self.current_k == self.k_max:
                    proposal = self._propose_death()
                else:
                    if np.random.rand() < 0.5:
                        proposal = self._propose_birth()
                    else:
                        proposal = self._propose_death()
            
            elif move_type == 'within_model':
                proposal = self._propose_within_model()

            # 处理提议
            if proposal and 'acceptance_ratio' in proposal:
                move_name = proposal.get('type', 'N/A')
                if np.random.rand() < proposal['acceptance_ratio']:
                    # 接受提议，更新状态
                    if 'k_new' in proposal: self.current_k = proposal['k_new']
                    if 'new_eigs' in proposal: self.current_eigenvalues = proposal['new_eigs']
                    self.current_a = self._cal_current_a_coeffs(self.current_eigenvalues)
                    if 'b_new' in proposal: self.current_b = proposal['b_new']
                    # sigma2_w 和 sigma2_z 在这些提议中没有更新，保持不变
                    accepted = True
            
            # 存储当前状态 (无论是否接受，都存储当前链的状态)
            k_history.append(self.current_k)
            a_history.append(self.current_a)
            b_history.append(self.current_b)
            sigma2_w_history.append(self.current_sigma2_w)
            sigma2_z_history.append(self.current_sigma2_z)
            eigenvalues_history.append(self.current_eigenvalues)

            if (i + 1) % 100 == 0:
                status = "Accepted" if accepted else "Rejected"
                print(f"Iteration {i+1}/{num_iterations}, k: {self.current_k}, Move: {move_name}, Status: {status}")

        return {
            "k": np.array(k_history),
            "a": a_history,
            "b": b_history,
            "sigma2_w": np.array(sigma2_w_history),
            "sigma2_z": np.array(sigma2_z_history),
            "eigenvalues": eigenvalues_history
        }
    

if __name__ == "__main__":
    # 导入生成数据的模块
    from generateSimData import simulate_lti_data
    from generateGroudTruth import generate_ground_truth_system
    import matplotlib.pyplot as plt
    
    # 配置 matplotlib 支持中文显示
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    # 1. 生成 Ground Truth 系统和仿真数据
    print("--- 1. 生成仿真数据 ---")
    # 真实系统阶数为 2
    A_true, B_true, C_true, D_true = generate_ground_truth_system(dx=2, rng_seed=42)
    
    T_steps = 400
    sigma_proc = 0.1  # 过程噪声标准差
    sigma_meas = 0.1  # 测量噪声标准差
    
    u_data, y_data = simulate_lti_data(
        A_true, B_true, C_true, D_true,
        T=T_steps,
        sigma_process=sigma_proc,
        sigma_measurement=sigma_meas,
        rng_seed=123
    )
    # RJMCMC_Sampler 需要一维的 y 和 u
    y_data_flat = y_data.flatten()
    u_data_flat = u_data.flatten()

    print(f"仿真数据生成完毕。 y shape: {y_data_flat.shape}, u shape: {u_data_flat.shape}")

    # 2. 初始化 RJMCMC 采样器
    print("\n--- 2. 初始化 RJMCMC 采样器 ---")
    sampler = RJMCMC_Sampler(
        y=y_data_flat,
        u=u_data_flat,
        k_min=1,
        k_max=4,      # 探索的最大阶数
        initial_k=3   # 从一个不等于真实阶数的阶数开始
    )
    
    # 3. 运行 MCMC
    print("\n--- 3. 开始运行 MCMC 采样 ---")
    num_iterations = 20000
    results = sampler.run_MCMC(num_iterations)
    print("MCMC 采样完成。")

    # 4. 分析和可视化结果
    print("\n--- 4. 分析结果 ---")
    
    # 丢弃早期样本 (burn-in)
    burn_in = 1000
    k_samples = results['k'][burn_in:]

    # 计算模型阶数的后验分布
    # 使用 bincount 统计每个 k 出现的次数
    k_posterior_counts = np.bincount(k_samples, minlength=sampler.k_max + 1)
    k_posterior_prob = k_posterior_counts / len(k_samples)

    # 打印后验概率
    print("模型阶数的后验概率分布:")
    for k_val in range(sampler.k_min, sampler.k_max + 1):
        print(f"  P(k={k_val} | y) ≈ {k_posterior_prob[k_val]:.4f}")

    # 可视化 k 的后验分布
    plt.figure(figsize=(10, 6))
    plt.bar(range(sampler.k_min, sampler.k_max + 1), 
            k_posterior_prob[sampler.k_min:sampler.k_max + 1], 
            color='skyblue', alpha=0.8, label='Posterior Probability')
    plt.axvline(x=2, color='red', linestyle='--', label='True Model Order (k=2)')
    plt.xlabel("模型阶数 (k)")
    plt.ylabel("后验概率 P(k|y)")
    plt.title("模型阶数的后验分布 (After Burn-in)")
    plt.xticks(range(sampler.k_min, sampler.k_max + 1))
    plt.legend()
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.show()