import numpy as np
from scipy.linalg import solve_discrete_lyapunov
import sys
import matplotlib.pyplot as plt
from generateGroudTruth import generate_ground_truth_system


def simulate_lti_data(A, B, C, D, T, sigma_process, sigma_measurement, rng_seed=None):
    """
    使用 LTI 系统 (A, B, C, D) 仿真生成数据。
    
    参数:
        A (np.ndarray): 状态矩阵 (dx x dx)
        B (np.ndarray): 输入矩阵 (dx x du)
        C (np.ndarray): 观测矩阵 (dy x dx)
        D (np.ndarray): 前馈矩阵 (dy x du)
        T (int): 轨迹长度 (时间步数)
        sigma_process (float): 过程噪声的标准差 (sigma_Sigma)
        sigma_measurement (float): 测量噪声的标准差 (sigma_Gamma)
        rng_seed (int, optional): 用于复现的随机种子

    返回:
        tuple: (u_data, y_data)
            u_data (np.ndarray): 输入轨迹 (T x du)
            y_data (np.ndarray): 输出轨迹 (T x dy)
    """
    
    # 初始化随机数生成器
    if rng_seed is None:
        rng = np.random.default_rng()
    else:
        rng = np.random.default_rng(rng_seed)
        
    # 从矩阵形状获取维度
    dx = A.shape[0]
    du = B.shape[1]
    dy = C.shape[0]
    
    # 初始化状态向量 x_0 = 0
    x = np.zeros((dx, 1))
    
    # 初始化用于存储历史数据的列表
    x_history = []
    y_history = []
    u_history = []
    
    # 生成噪声序列
    # 输入 u_t ~ N(0, I) 
    u_data_gen = rng.standard_normal(size=(T, du, 1))
    
    # 过程噪声 w_t ~ N(0, sigma_process^2 * I) [cite: 523]
    w_data_gen = rng.normal(scale=sigma_process, size=(T, dx, 1))
    
    # 测量噪声 z_t ~ N(0, sigma_measurement^2 * I) [cite: 523]
    z_data_gen = rng.normal(scale=sigma_measurement, size=(T, dy, 1))

    print(f"\n--- 开始仿真数据 (T={T}) ---")
    print(f"过程噪声 (sigma_Sigma): {sigma_process}")
    print(f"测量噪声 (sigma_Gamma): {sigma_measurement}")

    # 循环 T 个时间步
    for t in range(T):
        u = u_data_gen[t]
        w = w_data_gen[t]
        z = z_data_gen[t]
        
        # 1. 计算当前输出 y_t = C*x_t + D*u_t + z_t 
        y = C @ x + D @ u + z
        
        # 2. 计算下一个状态 x_{t+1} = A*x_t + B*u_t + w_t 
        x_next = A @ x + B @ u + w
        
        # 存储数据
        u_history.append(u.squeeze())
        y_history.append(y.squeeze())
        x_history.append(x.squeeze())
        
        # 更新状态
        x = x_next
        
    print("仿真完成。")
    
    # 将列表转换为 numpy 数组
    # 我们需要 (T, du) 和 (T, dy) 的形状
    # 使用 .reshape(-1, du) 和 .reshape(-1, dy) 来处理 du/dy=1 的情况
    u_data = np.array(u_history).reshape(T, du)
    y_data = np.array(y_history).reshape(T, dy)
    
    return u_data, y_data

if __name__ == '__main__':
    # --- 运行示例 ---
    
    # 1. 生成 Ground Truth 系统
    # 使用与第一步相同的种子，确保系统一致
    A_true, B_true, C_true, D_true = generate_ground_truth_system(rng_seed=42)
    
    # 2. 仿真数据
    # 根据 6.3 节的设置 
    T_steps = 400
    sigma_proc = 0.3
    sigma_meas = 0.0
    
    # 使用不同的种子进行仿真，以确保数据和系统生成是独立的
    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
    )
    
    # 打印结果形状
    print(f"\n--- 仿真结果 ---")
    print(f"输入数据 u_data 形状: {u_data.shape}")
    print(f"输出数据 y_data 形状: {y_data.shape}")

    # 可视化检查
    plt.figure(figsize=(12, 4))
    plt.plot(y_data, label=f'Simulated Output $y_t$ (noise $\sigma_w={sigma_proc}$)')
    plt.title('Simulated Data Trajectory (First Output Channel)')
    plt.xlabel('Time Step $t$')
    plt.ylabel('Output $y_t$')
    plt.legend()
    plt.grid(True)
    plt.show()