from typing import Any, List, Literal, Optional, Tuple, Union, cast

import numpy as np
import pywt
from scipy.signal import ShortTimeFFT
from scipy.signal.windows import get_window

from zyplib.utils.ensure import ensure_npy

__all__ = ['stft', 'istft', 'spectrogram', 'cwt', 'dwt']


def stft(
    x,
    fs,
    window='hamming',
    nperseg=None,
    noverlap=None,
    nfft=None,
    axis=-1,
    scaling='spectrum',
    return_onesided=True,
    **kwargs,
):
    """Short-Time Fourier Transform (STFT)

    计算短时傅里叶变换，返回复数频谱，所以通常更推荐使用 spectrogram 函数

    Parameters
    ----------
    - `x` : `array_like`
        - 输入信号。支持多维数组，沿指定轴进行变换。
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `window` : `str` or `tuple` or `array_like`, optional
        - 窗函数类型或窗函数数组。默认为 'hamming'。
        - 支持 scipy.signal.windows 中的所有窗函数。
    - `nperseg` : `int`, optional
        - 每段的长度。默认为 `int(0.5 * fs)`，即 0.5 秒。
    - `noverlap` : `int`, optional
        - 段间重叠的样本数。默认为 `nperseg // 2`。
    - `nfft` : `int`, optional
        - FFT 的长度。如果大于 `nperseg`，则进行零填充。默认等于 `nperseg`。
    - `axis` : `int`, optional
        - 进行变换的轴。默认为 -1（最后一个轴）。
    - `scaling` : `str`, optional
        - 缩放模式。'spectrum': 功率谱（V²）；'psd': 功率谱密度（V²/Hz）。默认为 'spectrum'。
        - 'spectrum' 返回功率谱，其模的平方是信号的功率谱，单位为 V²。
        - 'psd' 返回功率谱密度，其模的平方是信号的功率谱密度，单位为 V²/Hz。
    - `return_onesided` : `bool`, optional
        - 是否返回单边谱。默认为 True。

    Returns
    ----------
    - `f` : `ndarray`
        - 频率数组，单位为 Hz。
    - `t` : `ndarray`
        - 时间数组，单位为秒。
    - `Zxx` : `ndarray`
        - STFT 结果，复数数组。形状为 (..., freq_bins, time_bins)。

    Examples
    ----------
    ```py
    >>> import numpy as np
    >>> from zyplib.signal.time_frequency import stft
    >>>
    >>> # 生成测试信号
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> x = np.sin(2*np.pi*50*t) + 0.5*np.sin(2*np.pi*120*t)
    >>>
    >>> # 计算 STFT
    >>> f, t_stft, Zxx = stft(x, fs)
    >>> print(f"频率范围: {f[0]:.1f} - {f[-1]:.1f} Hz")
    >>> print(f"时间范围: {t_stft[0]:.2f} - {t_stft[-1]:.2f} s")
    ```
    """
    x = cast(np.ndarray, ensure_npy(x))

    # 设置默认参数
    if nperseg is None:
        nperseg = min(int(0.25 * fs), 1024)  # 更保守的默认值

    if noverlap is None:
        noverlap = nperseg // 2

    if nfft is None:
        nfft = nperseg

    # 计算 hop（步长）
    hop = nperseg - noverlap

    # 获取窗函数
    if isinstance(window, str) or isinstance(window, tuple):
        win = get_window(window, nperseg)
    else:
        win = np.asarray(window)
        if len(win) != nperseg:
            raise ValueError(f'窗长度 ({len(win)}) 必须等于 nperseg ({nperseg})')

    # 确定 FFT 模式
    fft_mode = 'onesided' if return_onesided else 'twosided'

    # 设置缩放模式
    scale_to = 'psd' if scaling == 'psd' else None

    # 创建 ShortTimeFFT 对象
    SFT = ShortTimeFFT(
        win=win, hop=hop, fs=fs, fft_mode=fft_mode, mfft=nfft, scale_to=scale_to, **kwargs
    )

    # 执行 STFT
    Zxx = SFT.stft(x, axis=axis, padding='even')  # 使用 even padding 缓解边缘效应

    # 获取频率和时间数组
    f = SFT.f
    n_samples = x.shape[axis]
    t = SFT.t(n_samples)

    return f, t, Zxx


def istft(
    Zxx,
    fs,
    nperseg,
    window='hamming',
    noverlap=None,
    nfft=None,
    time_axis=-1,
    freq_axis=-2,
    **kwargs,
):
    """Inverse Short-Time Fourier Transform (ISTFT)

    从 STFT 结果重构时域信号。

    Parameters
    ----------
    - `Zxx` : `array_like`
        - STFT 结果的复数数组。
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `nperseg` : `int`
        - 每段的长度，必须与计算 STFT 时使用的值相同。
    - `window` : `str` or `tuple` or `array_like`, optional
        - 窗函数类型，应与计算 STFT 时使用的窗函数相同。默认为 'hamming'。
    - `noverlap` : `int`, optional
        - 段间重叠的样本数，应与计算 STFT 时使用的值相同。默认为 `nperseg // 2`。
    - `nfft` : `int`, optional
        - FFT 的长度，应与计算 STFT 时使用的值相同。默认为 `nperseg`。
    - `time_axis` : `int`, optional
        - 时间轴的索引。默认为 -1。
    - `freq_axis` : `int`, optional
        - 频率轴的索引。默认为 -2。

    Returns
    ----------
    - `t` : `ndarray`
        - 重构信号的时间数组。
    - `x` : `ndarray`
        - 重构的时域信号。

    Notes
    -----
    为了确保重构的准确性，必须提供与 STFT 计算时完全相同的参数。
    特别是 `nperseg`、`noverlap` 和 `nfft` 必须与 `stft()` 调用时使用的值一致。
    """
    import warnings

    Zxx = ensure_npy(Zxx)

    # 获取 STFT 矩阵的形状信息
    freq_bins = Zxx.shape[freq_axis]

    # 验证参数
    if noverlap is None:
        noverlap = nperseg // 2
    elif noverlap >= nperseg:
        raise ValueError('noverlap 必须小于 nperseg')

    if nfft is None:
        nfft = nperseg
    elif nfft < nperseg:
        warnings.warn(
            'nfft < nperseg 可能导致重构不准确。请确保与 STFT 计算时的参数一致。',
            UserWarning,
            stacklevel=2,
        )

    # 计算 hop
    hop = nperseg - noverlap

    # 获取窗函数
    if isinstance(window, str) or isinstance(window, tuple):
        win = get_window(window, nperseg)
    else:
        win = np.asarray(window)

    # 确定 FFT 模式
    fft_mode = 'onesided' if freq_bins == nfft // 2 + 1 else 'twosided'

    # 创建 ShortTimeFFT 对象
    SFT = ShortTimeFFT(win=win, hop=hop, fs=fs, fft_mode=fft_mode, mfft=nfft, **kwargs)

    # 执行 ISTFT
    # 让 ShortTimeFFT 自动确定重构范围
    x = SFT.istft(Zxx, f_axis=freq_axis, t_axis=time_axis)

    # 计算时间数组
    n_samples = len(x) if x.ndim == 1 else x.shape[-1]
    t = np.arange(n_samples) / fs

    return t, x


def spectrogram(
    x,
    fs,
    window='hamming',
    nperseg=None,
    noverlap=None,
    nfft=None,
    axis=-1,
    scaling='density',
    mode='psd',
    **kwargs,
):
    """计算信号的谱图

    计算信号的功率谱图或幅度谱图; 返回 real 类型

    Parameters
    ----------
    - `x` : `array_like`
        - 输入信号。
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `window` : `str` or `tuple` or `array_like`, optional
        - 窗函数类型。默认为 'hamming'。
    - `nperseg` : `int`, optional
        - 每段的长度。默认为 `int(0.5 * fs)`。
    - `noverlap` : `int`, optional
        - 段间重叠的样本数。默认为 `nperseg // 2`。
    - `nfft` : `int`, optional
        - FFT 的长度。默认等于 `nperseg`。
    - `axis` : `int`, optional
        - 进行变换的轴。默认为 -1。
    - `scaling` : `str`, optional
        - 缩放模式。'density': 功率谱密度；'spectrum': 功率谱。默认为 'density'。
    - `mode` : `str`, optional
        - 输出模式。'magnitude': 幅度谱；'psd': 功率谱密度；'complex': 复数谱。默认为 'psd'。

    Returns
    ----------
    - `f` : `ndarray`
        - 频率数组，单位为 Hz。
    - `t` : `ndarray`
        - 时间数组，单位为秒。
    - `Sxx` : `ndarray`
        - 谱图。根据 `mode` 参数，可能是幅度谱、功率谱密度或复数谱。
    """
    # 计算 STFT
    f, t, Zxx = stft(
        x=x,
        fs=fs,
        window=window,
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft,
        axis=axis,
        scaling=scaling,
        **kwargs,
    )

    # 根据模式计算输出
    if mode == 'complex':
        Sxx = Zxx
    elif mode == 'magnitude':
        Sxx = np.abs(Zxx)
    elif mode == 'psd':
        Sxx = np.abs(Zxx) ** 2
    else:
        raise ValueError(
            f"不支持的模式: {mode}。支持的模式: 'magnitude', 'psd', 'complex'"
        )

    return f, t, Sxx


def cwt(
    x: np.ndarray,
    fs: float,
    wavelet: str = 'cmor1.5-1.0',
    scales: Optional[np.ndarray] = None,
    n_scales: int = 64,
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    axis: int = -1,
    method: Literal['fft', 'conv'] = 'fft',
):
    """连续小波变换 (Continuous Wavelet Transform, CWT)

    对信号进行连续小波变换，特别适用于低频成分重要的脑电信号分析。
    相比STFT，小波变换能提供更好的多分辨率分析能力：低频处有更高的频率分辨率，高频处有更高的时间分辨率。

    Parameters
    ----------
    - `x` : `np.ndarray`
        - 输入信号。支持多维数组，沿指定轴进行变换。
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `wavelet` : `str`, optional
        - 小波基函数名称，默认为 'morl'（Morlet小波）。
        - 支持 PyWavelets 库中的所有连续小波函数，包括：
          'cgau1'-'cgau8', 'cmor', 'fbsp', 'gaus1'-'gaus8', 'mexh', 'morl', 'shan'
    - `scales` : `np.ndarray`, optional
        - 小波变换的尺度参数数组。如果为None，根据n_scales, fmin和fmax自动生成
    - `n_scales` : `int`, optional
        - 自动生成的尺度数量，默认为64
    - `fmin` : `float`, optional
        - 最小频率，单位为Hz。默认为None，将被设置为fs/100（即较低频率）
    - `fmax` : `float`, optional
        - 最大频率，单位为Hz。默认为None，将被设置为fs/2（奈奎斯特频率）
    - `axis` : `int`, optional
        - 进行变换的轴。默认为-1（最后一个轴）。
    - `method` : `str`, optional
        - CWT计算方法，默认为 'fft'（FFT方法，速度更快）
        - 'conv': 使用卷积方法（适合短信号）
        - 'fft': 使用FFT方法（适合长信号，速度更快）

    Returns
    ----------
    - `scales` : `np.ndarray`
        - 尺度参数数组
    - `freqs` : `np.ndarray`
        - 对应的频率数组，单位为Hz
    - `t` : `np.ndarray`
        - 时间数组，单位为秒
    - `cwt_data` : `np.ndarray`
        - CWT结果，复数数组。形状为(..., n_scales, n_samples)。

    Examples
    ----------
    ```py
    >>> import numpy as np
    >>> from zyplib.signal.time_frequency import cwt
    >>>
    >>> # 生成测试信号
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> x = np.sin(2*np.pi*10*t) + np.sin(2*np.pi*50*t[:len(t)//2])
    >>>
    >>> # 计算CWT
    >>> scales, freqs, t_cwt, coef = cwt(x, fs, wavelet='cmor', fmin=1, fmax=100)
    >>> print(f"频率范围: {freqs[0]:.1f} - {freqs[-1]:.1f} Hz")
    >>> print(f"时间范围: {t_cwt[0]:.2f} - {t_cwt[-1]:.2f} s")
    ```

    Notes
    -----
    **EEG/MEG 脑电信号分析实践:**

    参数建议: 不设置 `scales`，而是通过设定感兴趣的**频率范围**来让函数自动计算最佳尺度。

    1.  **小波选择 (`wavelet`)**:
        - **Complex Morlet (`cmorB-C`, 如 'cmor1.5-1.0')** 是最常用的选择。它的复数特性使其能同时提供幅度和相位信息，非常适合计算时频功率和锁相值 (phase-locking value)。
        - `B` 是带宽参数，`C` 是中心频率参数。例如，`'cmor1.5-1.0'` 表示带宽为1.5，中心频率为1.0。

    2.  **频率范围 (`fmin`, `fmax`)**:
        - 根据您感兴趣的脑电频带进行设置。
        - **Delta (δ)**: 1-4 Hz
        - **Theta (θ)**: 4-8 Hz
        - **Alpha (α)**: 8-13 Hz
        - **Beta (β)**: 13-30 Hz
        - **Gamma (γ)**: 30-100 Hz
        - 若要进行全面的探索性分析，可设置 `fmin=1`, `fmax=100`。

    3.  **尺度数量 (`n_scales`)**:
        - 该参数控制频率分辨率。`64` 到 `256` 是常用范围。
        - 数量越多，时频图越平滑，但计算成本越高。对于发表级别的图像，通常选择 `128` 或更高。
        - 这个值决定了在`fmin`和`fmax`之间采样的频率点数量。值越大，频率轴越平滑，时频图越精细，但计算也越慢

    4.  **结果解读**:
        - 返回的 `cwt_data` 是复数。
        - **时频功率 (Time-Frequency Power)**: `power = np.abs(cwt_data)**2`。这是最常见的分析，用于观察特定频带能量随时间的变化（即事件相关同步/去同步, ERS/ERD）。
        - **相位 (Phase)**: `phase = np.angle(cwt_data)`。用于分析跨试次或跨通道的相位一致性。
    """
    x = cast(np.ndarray, ensure_npy(x))

    # 将输入数据移动到指定轴
    if axis != -1:
        x = np.moveaxis(x, axis, -1)

    # 计算信号长度
    n_samples = x.shape[-1]

    # 设置默认频率范围
    if fmin is None:
        fmin = fs / 100  # 较低频率，有助于捕捉低频成分
    if fmax is None:
        fmax = fs / 2  # 奈奎斯特频率

    # 确保频率范围有效
    fmin = max(fmin, fs / n_samples)  # 最低频率受信号长度限制
    fmax = min(fmax, fs / 2)  # 最高频率不超过奈奎斯特频率

    # 生成尺度参数
    if scales is None:
        # 将频率转换为尺度（频率与尺度成反比关系）
        # 获取中心频率，用于将尺度转换为频率
        wavelet_obj = cast(Any, pywt).ContinuousWavelet(wavelet)
        central_freq = cast(float, pywt.central_frequency(wavelet_obj))

        # 对数均匀尺度，确保在对数频率轴上均匀分布
        scales = central_freq * fs / np.logspace(np.log10(fmin), np.log10(fmax), n_scales)
        scales = cast(np.ndarray, scales)[::-1]  # 反转顺序，使低频（大尺度）在上方

    # 计算小波变换并获取频率信息
    cwt_data, freqs = pywt.cwt(
        x,
        scales=scales,
        wavelet=wavelet,
        sampling_period=1.0 / fs,
        method=method,
        axis=-1,
    )

    # 确保频率按降序排列（高频在前，与尺度对应）
    if len(freqs) > 1 and freqs[0] < freqs[-1]:
        freqs = freqs[::-1]
        cwt_data = cwt_data[..., ::-1, :]

    # 将结果移回原始轴
    if axis != -1:
        cwt_data = np.moveaxis(cwt_data, -1, axis)

    # 生成时间数组
    t = np.arange(0, n_samples) / fs

    return scales, freqs, t, cwt_data


def dwt(
    x: np.ndarray,
    fs: float,
    wavelet: str = 'db4',
    level: Optional[int] = None,
    mode: str = 'symmetric',
    axis: int = -1,
):
    """离散小波变换 (Discrete Wavelet Transform, DWT)

    对信号进行多分辨率分析，分解为近似系数和多个尺度的细节系数。
    适用于提取信号的时频特征，特别是脑电信号中的不同频段活动。

    Parameters
    ----------
    - `x` : `np.ndarray`
        - 输入信号。支持多维数组，沿指定轴进行变换。
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `wavelet` : `str`, optional
        - 小波基函数名称，默认为'db4'（Daubechies 4阶小波）。
        - 常用选择：'db'系列(db1-db38), 'sym'系列, 'coif'系列, 'haar'等
    - `level` : `int`, optional
        - 分解层数。如果为None，则自动选择最大可分解层数。
    - `mode` : `str`, optional
        - 边界处理模式，默认为'symmetric'（对称扩展）。
    - `axis` : `int`, optional
        - 进行变换的轴。默认为-1（最后一个轴）。

    Returns
    ----------
    - `coeffs` : `list`
        - 小波系数列表。[cAn, cDn, cDn-1, ..., cD1]，其中cAn是最后一层的近似系数，cDi是第i层的细节系数。
    - `freqs` : `list`
        - 各层对应的频率带列表，每个元素为(低频界, 高频界)元组，单位为Hz。
    - `t` : `list`
        - 各层对应的时间数组列表，单位为秒。

    Examples
    ----------
    ```py
    >>> import numpy as np
    >>> from zyplib.signal.time_frequency import dwt
    >>>
    >>> # 生成测试信号
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> x = np.sin(2*np.pi*10*t) + 0.5*np.sin(2*np.pi*30*t) + 0.2*np.random.randn(len(t))
    >>>
    >>> # 计算DWT
    >>> coeffs, freqs, t_list = dwt(x, fs, wavelet='db4', level=5)
    >>>
    >>> # 访问第3层的细节系数(对应频带约为62.5-125 Hz)
    >>> d3 = coeffs[-3]
    >>> freq_band = freqs[-3]
    >>> t3 = t_list[-3]
    >>> print(f"第3层细节系数频带: {freq_band[0]:.1f}-{freq_band[1]:.1f} Hz")
    ```

    Notes
    -----
    **EEG/MEG 脑电信号分析实践:**

    DWT 非常适合从EEG信号中**提取频带特征**，用于机器学习模型（如睡眠分期、癫痫检测、BCI等）。

    1.  **小波选择 (`wavelet`)**:
        - **Daubechies (`'db4'`, `'db6'`)** 或 **Symlets (`'sym4'`, `'sym8'`)** 是常用选择
        - 常用推荐: `'db4'` 或 `'sym4'`

    2.  **分解层数 (`level`)**:
        - **这是最关键的参数，必须根据您的采样率 `fs` 来确定。**
        - **计算方法**: 对于第 `k` 层的细节系数 `cDk`，其频带约为 `[fs / 2**(k+1), fs / 2**k]`。
        - **示例 (fs=250 Hz)**:
          - `level=6` 的分解：
            - `cD1`: [62.5, 125] Hz (高Gamma)
            - `cD2`: [31.25, 62.5] Hz (Gamma)
            - `cD3`: [15.625, 31.25] Hz (Beta)
            - `cD4`: [7.81, 15.625] Hz (Alpha)
            - `cD5`: [3.91, 7.81] Hz (Theta)
            - `cD6`: [1.95, 3.91] Hz (Delta)
            - `cA6`: [0, 1.95] Hz (低Delta)
        - 您需要根据您的 `fs` 预先计算好这个对应表。

    3.  **特征提取**:
        - 分解后，`coeffs` 列表包含了每个频带的系数。
        - 常用的特征包括：
          - **能量 (Energy)**: `np.sum(np.square(cD_k))`
          - **功率 (Power)**: `np.mean(np.square(cD_k))`
          - **标准差 (Standard Deviation)**: `np.std(cD_k)`
          - **熵 (Entropy)**: `scipy.stats.entropy(np.square(cD_k))`
        - 将这些特征组合起来，可以构成用于分类器的特征向量。
    """
    x = cast(np.ndarray, ensure_npy(x))

    # 将输入数据移动到指定轴
    if axis != -1:
        x = np.moveaxis(x, axis, -1)

    # 计算信号长度
    n_samples = x.shape[-1]

    # 自动确定最大分解层数
    max_level = pywt.dwt_max_level(
        data_len=n_samples, filter_len=cast(Any, pywt).Wavelet(wavelet).dec_len
    )
    if level is None:
        level = max_level
    else:
        level = min(level, max_level)  # 不超过最大可分解层数
    level = cast(int, level)

    # 计算小波分解
    coeffs = pywt.wavedec(x, wavelet=wavelet, level=level, mode=mode, axis=-1)

    # 计算各层的频率带
    freqs = []
    current_fs = fs

    # 对于近似系数（最低频带）
    low_freq = 0
    high_freq = current_fs / 2 ** (level + 1)
    freqs.append((low_freq, high_freq))

    # 对于各层细节系数（从最深层到浅层）
    for i in range(level):
        # 计算当前层的频率带
        low_freq = current_fs / 2 ** (level - i + 1)
        high_freq = current_fs / 2 ** (level - i)
        freqs.append((low_freq, high_freq))

    # 计算各层的时间数组
    t_list = []
    for i, coef in enumerate(coeffs):
        # 获取当前层系数的长度
        coef_len = coef.shape[-1]
        # 计算时间步长（较深层的时间分辨率更低）
        # 计算每层的下采样因子
        downsample_factor = 2 ** (level - i + 1) if i > 0 else 2**level
        # 计算时间步长
        dt = downsample_factor / fs
        t_layer = np.arange(coef_len) * dt
        t_list.append(t_layer)

    return coeffs, freqs, t_list
