import time
import numpy as np
from matplotlib import pyplot as plt
from reservoirpy import datasets

def print_exec_time(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        print(f"Function {func.__name__} executed in {time.time() - start:.4f} seconds.")
        return result
    return wrapper

def add_noise(data, noise_type='gaussian', intensity=0.1, **kwargs):
    """
    为输入数据添加噪声。

    参数:
        data: numpy 数组, 原始数据.
        noise_type: str, 噪声类型，可选 'gaussian'（高斯白噪声）、'colored' 或 '1/f'（色噪声）、'impulse'（脉冲噪声）、'sine'（正弦噪声）。
        intensity: float, 噪声强度.
        kwargs: 针对某些噪声类型的额外参数，例如:
                - 对于 'impulse': prob (脉冲发生概率, 默认 0.01)
                - 对于 'sine'   : freq (正弦信号频率, 默认 5)
    返回:
        添加噪声后的数据.
    """
    noise_type = noise_type.lower()
    
    # 确保噪声与数据形状匹配，处理多维数据
    if data.ndim > 1:
        n_samples, n_dimensions = data.shape
    else:
        n_samples = len(data)
        n_dimensions = 1
    
    if noise_type in ['gaussian', 'gaussian white']:
        noise = intensity * np.random.randn(*data.shape)
        return data + noise

    elif noise_type in ['colored', '1/f']:
        # 对于多维数据，为每个维度单独生成1/f噪声
        if data.ndim > 1:
            noise = np.zeros_like(data)
            for dim in range(n_dimensions):
                f = np.fft.fftfreq(n_samples)
                f[0] = 1e-10
                noise_complex = np.random.randn(n_samples) + 1j * np.random.randn(n_samples)
                factor = intensity / np.sqrt(np.abs(f))
                noise[:, dim] = np.fft.ifft(noise_complex * factor).real
        else:
            f = np.fft.fftfreq(n_samples)
            f[0] = 1e-10
            noise_complex = np.random.randn(n_samples) + 1j * np.random.randn(n_samples)
            factor = intensity / np.sqrt(np.abs(f))
            noise = np.fft.ifft(noise_complex * factor).real
        return data + noise

    elif noise_type == 'impulse':
        noise = np.zeros_like(data)
        prob = kwargs.get('prob', 0.01)
        mask = np.random.rand(*data.shape) < prob
        impulse = intensity * (2 * np.random.rand(*data.shape) - 1)
        noise[mask] = impulse[mask]
        return data + noise

    elif noise_type == 'sine':
        t = np.arange(n_samples)
        freq = kwargs.get('freq', 5)
        
        # 对于多维数据，确保正弦波形状正确
        if data.ndim > 1:
            sine_wave = intensity * np.sin(2 * np.pi * freq * t / n_samples).reshape(-1, 1)
            # 扩展到与数据相同的维度
            sine_wave = np.tile(sine_wave, (1, n_dimensions))
        else:
            sine_wave = intensity * np.sin(2 * np.pi * freq * t / n_samples)
        return data + sine_wave

    else:
        raise ValueError("未知的噪声类型。可选项：'gaussian', 'colored' (1/f), 'impulse', 'sine'")

def load_data(system='lorenz', init='random', noise=None, intensity=0.1, h=0.01, n_timesteps=10000, transient=1000, normlization=True, **kwargs):
    """
    加载混沌系统数据.

    参数:
        system: str, 混沌系统类型, 可选 'lorenz', 'rossler', 'multiscroll', 'kuramoto_sivashinsky'。
        init: 初始值. 如果为 'random' 则随机初始化，否则期望传入数组形式的初始值.
        noise: None, str 或 str 列表, 指定要添加的噪声类型. 如果是列表，则混合各噪声 (取平均) 作为混合噪声.
        -- 可选 'gaussian'（高斯白噪声）、'colored' 或 '1/f'（色噪声）、'impulse'（脉冲噪声）、'sine'（正弦噪声）
        intensity: float, 噪声强度.
        h: float, 时间步长 (TimeDelta).
        n_timesteps: int, 时间步数.
        transient: int, 需丢弃的时间步数.
        normlization: bool, 是否对数据进行归一化处理.
        kwargs: 其他系统参数.
    返回:
        (clean_data, noisy_data): 两个 numpy 数组, 分别为干净数据和添加噪声后的数据.
        
    各系统默认值:
        - lorenz: rho=28, sigma=10, beta=8/3, h=0.03, x0=[1, 1, 1]
        - rossler: a=0.2, b=0.2, c=5.7, h=0.1, x0=[-0.1, 0, 0.02]
        - multiscroll: a=40, b=3, c=28, h=0.01, x0=[-0.1, 0.5, -0.6]
        - kuramoto_sivashinsky: N=128, M=16, h=0.25, x0=None
    """
    system = system.lower()
    if init == 'random':
        if system in ['lorenz', 'rossler', 'multiscroll']:
            # 默认3维初始值
            state = np.random.rand(3)
        elif system in ['kuramoto_sivashinsky']:
            state = np.random.rand(1)
        else:
            raise ValueError("未知的混沌系统类型。")
    else:
        state = np.array(init, dtype=float)

    if system == 'lorenz':
        '''
        (function) def lorenz(
            n_timesteps: int,
            rho: float = 28,
            sigma: float = 10,
            beta: float = 8 / 3,
            x0: list | ndarray = [1, 1, 1],
            h: float = 0.03,
            **kwargs: Any
        ) -> ndarray
        '''
        sigma = kwargs.get('sigma', 10.0)
        rho   = kwargs.get('rho', 28.0)
        beta  = kwargs.get('beta', 8/3)
        clean_data = datasets.lorenz(n_timesteps=n_timesteps, h=h, sigma=sigma, rho=rho, beta=beta, x0=state)

    elif system == 'rossler':
        '''
        (function) def rossler(
            n_timesteps: int,
            a: float = 0.2,
            b: float = 0.2,
            c: float = 5.7,
            x0: list | ndarray = [-0.1, 0, 0.02],
            h: float = 0.1,
            **kwargs: Any
        ) -> ndarray
        '''
        a = kwargs.get('a', 0.2)
        b = kwargs.get('b', 0.2)
        c = kwargs.get('c', 5.7)
        clean_data = datasets.rossler(n_timesteps=n_timesteps, h=h, a=a, b=b, c=c, x0=state)

    elif system == 'multiscroll':
        '''
        (function) def multiscroll(
            n_timesteps: int,
            a: float = 40,
            b: float = 3,
            c: float = 28,
            x0: list | ndarray = [-0.1, 0.5, -0.6],
            h: float = 0.01,
            **kwargs: Any
        ) -> ndarray
        '''
        a = kwargs.get('a', 40)
        b = kwargs.get('b', 3)
        c = kwargs.get('c', 28)
        clean_data = datasets.multiscroll(n_timesteps=n_timesteps, h=h, x0=state, a=a, b=b, c=c)

    elif system == 'kuramoto_sivashinsky':
        '''
        (function) def kuramoto_sivashinsky(
            n_timesteps: int,
            warmup: int = 0,
            N: int = 128,
            M: float = 16,
            x0: list | ndarray = None,
            h: float = 0.25
        ) -> ndarray
        '''
        clean_data = datasets.kuramoto_sivashinsky(n_timesteps=n_timesteps, h=h, **kwargs)

    else:
        raise ValueError("未知的混沌系统类型。可选项: 'lorenz', 'rossler', 'multiscroll', 'kuramoto_sivashinsky'")

    if normlization:
        clean_data = (clean_data - clean_data.mean(axis=0)) / clean_data.std(axis=0) # Z-score 归一化
    clean_data = clean_data[transient:]
    
    # 添加噪声
    if noise is not None:
        if isinstance(noise, list):
            noise_sum = np.zeros_like(clean_data)
            for nt in noise:
                noise_sum += add_noise(np.zeros_like(clean_data), noise_type=nt, intensity=intensity, **kwargs)
            mixed_noise = noise_sum / len(noise)
            noisy_data = clean_data + mixed_noise
        elif isinstance(noise, str):
            noisy_data = add_noise(clean_data, noise_type=noise, intensity=intensity, **kwargs)
        else:
            raise ValueError("噪声参数 noise 必须为字符串或字符串列表。")
    else:
        noisy_data = clean_data.copy()

    return clean_data, noisy_data

def visualize_data_diff(clean_data, noisy_data, prediction, warmup=0, title_prefix=""):
    plt.figure(figsize=(12, 6))
    for i in range(3):
        plt.subplot(3, 2, 2*i+1)
        plt.plot(noisy_data[warmup:, i], label='noisy')
        plt.plot(clean_data[warmup:, i], label='clean')
        plt.title(f"{title_prefix} Dim {i+1} (Noisy vs Clean)")
        plt.legend()
        plt.subplot(3, 2, 2*i+2)
        plt.plot(prediction[warmup:, i], label='prediction')
        plt.plot(clean_data[warmup:, i], label='clean')
        plt.title(f"{title_prefix} Dim {i+1} (Prediction vs Clean)")
        plt.legend()
    plt.tight_layout()

def visualize_psd_diff(clean_data, noisy_data, prediction, warmup=0, title_prefix=""):
    plt.figure(figsize=(12, 6))
    for i in range(3):
        plt.subplot(3, 1, i+1)
        plt.psd(clean_data[warmup:, i], NFFT=1024, label='clean')
        plt.psd(noisy_data[warmup:, i], NFFT=1024, label='noisy')
        plt.psd(prediction[warmup:, i], NFFT=1024, label='prediction')
        plt.title(f"{title_prefix} Dim {i+1} PSD Comparison")
        plt.legend()
    plt.tight_layout()

if __name__ == "__main__":
    # 测试数据加载和噪声添加
    clean_data, noisy_data = load_data(system='lorenz', noise=['gaussian', 'colored'], intensity=0.1, n_timesteps=10000, transient=1000)
    print("Clean Data Shape:", clean_data.shape)
    print("Noisy Data Shape:", noisy_data.shape)