import numpy as np
from scipy.linalg import solve_discrete_lyapunov
import sys

def generate_ground_truth_system(dx=2, du=1, dy=1, rng_seed=None):
    """
    生成一个 "真实" 的、稳定的、可控的、可观测的 LTI 系统。
    
    该过程遵循论文 6.2 节中描述的方法。
    
    参数:
        dx (int): 状态维度 (state dimension)
        du (int): 输入维度 (input dimension)
        dy (int): 输出维度 (output dimension)
        rng_seed (int, optional): 用于复现的随机种子
        
    返回:
        tuple: (A, B, C, D) 矩阵
    """
    
    # 初始化随机数生成器
    if rng_seed is None:
        rng = np.random.default_rng()
    else:
        rng = np.random.default_rng(rng_seed)

    # 尝试生成一个良态的系统，最多重试 100 次
    max_retries = 100
    for attempt in range(max_retries):
        try:
            # --- 1. 生成稳定的 A 矩阵 (dx x dx) ---
            # 论文图 2(b) 显示了复特征值，遵循 6.2 节的极坐标法
            
            # 在极坐标下采样一对共轭复特征值
            r_squared = rng.uniform(0, 1.0) # 采样 r^2，确保在单位圆内
            r = np.sqrt(r_squared)
            theta = rng.uniform(0, np.pi)   # 仅在上半平面采样角度
            
            lambda1 = r * (np.cos(theta) + 1j * np.sin(theta))
            lambda2 = np.conjugate(lambda1)
            
            # 创建对应的实数块对角矩阵
            alpha, beta = lambda1.real, lambda1.imag
            lambda_block = np.array([[alpha, beta], 
                                    [-beta, alpha]])
            
            # 生成一个随机正交矩阵 V (dx x dx) [cite: 511, 584]
            Z = rng.standard_normal(size=(dx, dx))
            V, _ = np.linalg.qr(Z)
            
            # 组装 A = V * Lambda_block * V^T
            A = V @ lambda_block @ V.T

            # --- 2. 生成 B (dx x du) 和 C (dy x dx) 矩阵 ---
            # 元素从 N(0, 1) 独立采样 
            B = rng.standard_normal(size=(dx, du))
            C = rng.standard_normal(size=(dy, dx))

            # --- 3. 检查可控性和可观测性  ---
            # 求解离散时间李雅普诺夫方程 (Lyapunov equation)
            Wc = solve_discrete_lyapunov(A, B @ B.T) # Controllability Gramian
            Wo = solve_discrete_lyapunov(A.T, C.T @ C) # Observability Gramian

            # 检查 Gramian 矩阵的条件
            # "拒绝主要特征值占总能量 99% 以上的系统"
            eig_Wc = np.linalg.eigvalsh(Wc)
            eig_Wo = np.linalg.eigvalsh(Wo)

            cond_c = np.max(eig_Wc) / np.sum(eig_Wc)
            cond_o = np.max(eig_Wo) / np.sum(eig_Wo)

            # 如果系统是良态的，则跳出循环
            if cond_c < 0.99 and cond_o < 0.99:
                
                # --- 4. 定义 D 矩阵 (dy x du) ---
                # 算例 6.3 中 D=0 
                D = np.zeros((dy, du))
                
                print(f"--- 成功生成 Ground Truth 系统 (尝试次数: {attempt + 1}) ---")
                print(f"特征值: {lambda1:.4f}, {lambda2:.4f}")
                print(f"Gramian 能量占比: Wc={cond_c:.4f}, Wo={cond_o:.4f}")
                print("A = \n", A)
                print("B = \n", B)
                print("C = \n", C)
                print("D = \n", D)
                
                return A, B, C, D
                
        except np.linalg.LinAlgError:
            # 李雅普诺夫方程求解器可能失败（例如，A 矩阵数值上不稳定）
            print(f"尝试 {attempt + 1} 失败 (LinAlgError)。正在重试...", file=sys.stderr)
            continue
            
    # 如果循环结束仍未成功
    raise RuntimeError(f"在 {max_retries} 次尝试后未能生成一个良态的系统。")

if __name__ == '__main__':
    # --- 运行示例 ---
    # 设置一个随机种子，以便每次运行时都能得到相同的结果
    A_true, B_true, C_true, D_true = generate_ground_truth_system(rng_seed=42)