import numpy as np
from numpy.typing import NDArray


def produce_wave_by_psd(
    f: NDArray, psd: NDArray, sfreq: int, t_second: int, n_sample: int = 1
) -> NDArray:
    """根据功率谱密度生成信号

    传入一个 PSD 的频谱图, 模拟生成一个拥有对应 PSD 的信号

    Parameters
    ----------
    - `f` : `NDArray`
        - 频率轴，单位为Hz
    - `psd` : `NDArray`
        - 输入的功率谱密度图
    - `sfreq` : `int`
        - 采样频率，单位为Hz
    - `t_second` : `int`
        - 信号时长，单位为秒
    - `n_sample` : `int`, optional
        - 生成的信号样本数量, by default 1

    Returns
    ----------
    - `NDArray`
        - 生成的信号，形状为 (n_sample, n_timepoints)
    """
    # 计算信号长度
    n_timepoints = int(sfreq * t_second)

    # 构造频率轴（对应FFT的频率bins）
    freqs = np.fft.fftfreq(n_timepoints, 1.0 / sfreq)

    # 只取正频率部分（包括0频率）
    positive_freqs = freqs[: n_timepoints // 2 + 1]

    # 将输入的PSD插值到FFT频率网格上
    psd_interp = np.interp(positive_freqs, f, psd)

    # 从PSD计算幅度谱 (PSD = |X(f)|^2, 所以 |X(f)| = sqrt(PSD))
    amplitude = np.sqrt(psd_interp)

    # 生成多个信号样本
    signals = np.zeros((n_sample, n_timepoints))

    for i in range(n_sample):
        # 生成随机相位 (0 to 2π)
        phases = np.random.uniform(0, 2 * np.pi, len(positive_freqs))

        # DC分量（0Hz）的相位应该为0（实数）
        phases[0] = 0

        # 如果信号长度为偶数，Nyquist频率的相位也应该为0
        if n_timepoints % 2 == 0:
            phases[-1] = 0

        # 构造复数频谱
        complex_spectrum = amplitude * np.exp(1j * phases)

        # 构造完整的频谱（包括负频率部分）
        # 负频率部分是正频率部分的共轭
        if n_timepoints % 2 == 0:
            # 偶数长度
            full_spectrum = np.concatenate(
                [complex_spectrum, np.conj(complex_spectrum[-2:0:-1])]
            )
        else:
            # 奇数长度
            full_spectrum = np.concatenate(
                [complex_spectrum, np.conj(complex_spectrum[-1:0:-1])]
            )

        # 逆FFT得到时域信号
        signal = np.fft.ifft(full_spectrum).real

        # 归一化信号以匹配PSD的能量
        # 计算信号的实际PSD
        signal_fft = np.fft.fft(signal)
        signal_psd = np.abs(signal_fft[: len(positive_freqs)]) ** 2 / (
            sfreq * n_timepoints
        )

        # 计算缩放因子
        target_power = np.trapz(psd_interp, positive_freqs)
        actual_power = np.trapz(signal_psd, positive_freqs)

        if actual_power > 0:
            scale_factor = np.sqrt(target_power / actual_power)
            signal *= scale_factor

        signals[i] = signal

    return signals if n_sample > 1 else signals[0]


def _simulate_eeg_psd(freqs: np.ndarray) -> np.ndarray:
    """创建类似EEG的功率谱密度 (PSD)

    该函数建模典型静息态EEG信号的PSD特征

    Parameters
    ----------
    - `freqs` : `np.ndarray`
        - 频率数组，单位为Hz

    Returns
    ----------
    - `np.ndarray`
        - EEG样式的功率谱密度
    """
    # 1/f background
    background = 10 / (freqs + 0.1) ** 1.2
    # Delta activity (0.5-4 Hz)
    delta_peak = 20 * np.exp(-(((freqs - 2) / 1) ** 2))
    # Theta activity (4-8 Hz)
    theta_peak = 8 * np.exp(-(((freqs - 6) / 2) ** 2))
    # Alpha peak (8-12 Hz)
    alpha_peak = 15 * np.exp(-(((freqs - 10) / 2) ** 2))
    # Beta activity (13-30 Hz)
    beta_peak = 5 * np.exp(-(((freqs - 20) / 8) ** 2))
    # High frequency rolloff
    rolloff = np.exp(-((freqs / 50) ** 2))

    psd = (background + delta_peak + alpha_peak + beta_peak + theta_peak) * rolloff
    return np.maximum(psd, 0.01)


def produce_fake_eeg_waves(
    duration: int,
    sfreq: int,
    n_signals: int = 1,
    psd_freq_range: tuple[float, float] = (0.2, 60),
    target_psd_resolution: float = 0.25,
) -> NDArray:
    """生成具有典型脑电图（EEG）特征的模拟信号

    Parameters
    ----------
    - `duration` : `int`
        - 信号时长 (秒)
    - `sfreq` : `int`
        - 采样频率 (Hz)
    - `n_signals` : `int`, optional
        - 生成的信号数量, by default 1
    - `psd_freq_range` : `tuple[float, float]`, optional
        - 用于构建目标PSD的频率范围, by default (0.5, 60)
    - `target_psd_resolution` : `int`, optional
        - 用于构建目标PSD的频率分辨率, by default 0.25

    Returns
    ----------
    - `a_signals (NDArray)`: 生成的信号, shape (n_signals, n_timepoints) 当 n_signals > 1 时，或 (n_timepoints,) 当 n_signals = 1 时
    """
    f_min, f_max = psd_freq_range
    num_points = int((f_max - f_min) / target_psd_resolution) + 1
    psd_freq_points = num_points if num_points > 1 else 2
    # 1. 创建目标PSD
    freqs = np.linspace(f_min, f_max, psd_freq_points)
    target_psd = _simulate_eeg_psd(freqs)

    # 2. 使用已有的 `produce_signal_by_psd` 函数生成信号
    a_signals = produce_wave_by_psd(freqs, target_psd, sfreq, duration, n_signals)

    return a_signals
