import numpy as np
from typing import cast
from scipy import signal

from zyplib.utils.ensure import ensure_npy

__all__ = [
    'welch',
    'psd_welch',
    'bandpower',
    'eeg_delta_power',
    'eeg_theta_power',
    'eeg_alpha_power',
    'eeg_sigma_power',
    'eeg_beta_power',
]


def welch(
    x,
    fs,
    nperseg: int | None = None,
    noverlap: int | None = None,
    nfft: int | None = None,
    max_freq=None,
    window='hamming',
    axis=-1,
    scaling='density',
    return_db=False,
    **scipy_welch_kwargs,
):
    """Welch PSD

    计算 welch 功率谱, 最后截断结果, 返回低于 lowpass Hz 的部分

    频率分辨率权衡：
    - 频率分辨率 = fs / nperseg
    - 更高的频率分辨率需要更长的窗口，但会降低时间分辨率和统计稳定性
    - 更长的窗口会减少分段数量，降低噪声抑制能力

    Parameters
    ----------
    - `x` : `array_like`
        - 输入信号。
    - `fs` : `int`
        - 采样率。
    - `nperseg` : `int`, optional
        - 每段的长度, 默认 `2 * fs`，频率分辨率为 `fs / nperseg` (默认 0.5 Hz)。
    - `noverlap` : `int`, optional
        - 段之间的重叠点数。如果为 `None`，默认为 `nperseg // 2`。
    - `nfft` : `int`, optional
        - FFT 长度。如果为 `None`，默认为大于等于 `nperseg` 的最小 2 的幂次。增加 `nfft` 只能插值频谱，不能真正提高频率分辨率。
    - `max_freq` : `int`, optional
        - 截断频率。当非 `None` 时, 截断超过 `max_freq` 的部分。
    - `axis` : `int`, optional
        - 默认 `-1`，即最后一个维度。
    - `scaling` : `str`, optional
        - 默认 `'density'`。`'density'`: 返回功率谱密度 (功率/Hz)；`'spectrum'`: 返回功率谱 (功率)。
    - `return_db` : `bool`, optional
        - 是否将功率谱转换为分贝 (dB) 形式。`True`: 返回 `10*log10(Pxx)`；`False`: 返回原始功率谱。

    Returns
    ----------
    - `f` : `ndarray`
        - 频率，分辨率为 `fs / nperseg`，长度为 `nperseg / 2`。
    - `Pxx` : `ndarray`
        - 功率谱或功率谱密度。当 `return_db=True` 时, 单位为 dB。
    """
    x = cast(np.ndarray, ensure_npy(x))
    if not nperseg:
        nperseg = fs * 2
    nperseg = cast(int, nperseg)
    if nfft is None:
        nfft = int(np.power(2, np.ceil(np.log2(nperseg))))
    f, Pxx = signal.welch(
        x,
        fs,
        window=window,
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft,
        axis=axis,
        scaling=scaling,
        **scipy_welch_kwargs,
    )
    # 当需要截断且截断频率小于 fs / 2 时, 进行截断
    if max_freq is not None and max_freq < fs / 2:
        valid_indices = np.where(f <= max_freq)[0]
        if len(valid_indices) == 0:
            raise ValueError(
                f'max_freq ({max_freq} Hz) is too low. Minimum frequency is {f[0] if len(f) > 0 else "N/A" } Hz.'
            )
        ind_cutoff = valid_indices[-1] + 1  # Slice up to this index
        f = f[:ind_cutoff]
        Pxx = Pxx[..., :ind_cutoff]  # Assumes Pxx's last axis is frequency

    # 转换为分贝形式 (如果需要)
    if return_db:
        Pxx = 10 * np.log10(Pxx)

    return f, Pxx


psd_welch = welch  # Alias for compatibility


def bandpower(
    x,
    fs,
    band: tuple,
    axis=-1,
    nperseg=None,
    scaling='density',
    return_db=False,
    **scipy_welch_kwargs,
):
    """计算特定频带的功率

    Parameters
    ----------
    - `x` : `array_like`
        - 输入信号。
    - `fs` : `int`
        - 采样率。
    - `band` : `tuple`
        - 频带范围 (`低频`, `高频`)。
    - `axis` : `int`, optional
        - 默认 `-1`，即最后一个维度。
    - `nperseg` : `int`, optional
        - 每段的长度, 默认 `2 * fs`。
    - `scaling` : `str`, optional
        - `'density'`: 返回功率谱密度 (功率/Hz)；`'spectrum'`: 返回功率谱 (功率)。
    - `return_db` : `bool`, optional
        - 是否将功率谱转换为分贝 (dB) 形式。

    Returns
    ----------
    - `power` : `ndarray`
        - 指定频带的功率。
    """
    x = ensure_npy(x)
    f, Pxx = welch(
        x,
        fs,
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )
    ind_min = np.argmax(f > band[0])
    ind_max = np.argmax(f > band[1])
    if ind_min == 0 or ind_max == 0:
        raise ValueError('band values are too low')
    return np.trapz(Pxx[..., ind_min:ind_max], f[ind_min:ind_max], axis=axis)


def eeg_delta_power(
    x, fs, axis=-1, nperseg=None, scaling='density', return_db=False, **scipy_welch_kwargs
):
    """Delta 波段功率 (0.5-4 Hz)

    计算 EEG 信号的 Delta 波段 (0.5-4 Hz) 功率。

    Parameters
    ----------
    - `x` : array_like, 1-D or 2-D
        - 待处理的原始 EEG 信号, 可以是一维或二维数组, 类型为 `array_like`。
    - `fs` : float
        - 输入信号的采样频率, 单位为 Hz, 类型为 `float`。
    - `axis` : int, optional
        - 可选参数, 计算取样信号频谱的轴。默认值为 -1, 即最后一个轴, 类型为 `int`。
    - `nperseg` : int, optional
        - 每段的长度, 默认 2 * fs
    - `scaling` : str, optional
        - 'density': 返回功率谱密度 (功率/Hz)
        - 'spectrum': 返回功率谱 (功率)
    - `return_db` : bool, optional
        - 是否将功率谱转换为分贝 (dB) 形式

    Returns
    ----------
    - `band_power` : `float`
        - Delta 波段功率。
    """

    return bandpower(
        x,
        fs,
        (0.5, 4),
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )


def eeg_theta_power(
    x, fs, axis=-1, nperseg=None, scaling='density', return_db=False, **scipy_welch_kwargs
):
    """Theta 波段功率 (4-8 Hz)

    计算 EEG 信号的 Theta 波段 (4-8 Hz) 功率。

    Parameters
    ----------
    - `x` : array_like, 1-D or 2-D
        - 待处理的原始 EEG 信号, 可以是一维或二维数组, 类型为 `array_like`。
    - `fs` : float
        - 输入信号的采样频率, 单位为 Hz, 类型为 `float`。
    - `axis` : int, optional
        - 可选参数, 计算取样信号频谱的轴。默认值为 -1, 即最后一个轴, 类型为 `int`。
    - `nperseg` : int, optional
        - 每段的长度, 默认 2 * fs
    - `scaling` : str, optional
        - 'density': 返回功率谱密度 (功率/Hz)
        - 'spectrum': 返回功率谱 (功率)
    - `return_db` : bool, optional
        - 是否将功率谱转换为分贝 (dB) 形式

    Returns
    ----------
    - `band_power` : `float`
        - Theta 波段功率。
    """
    return bandpower(
        x,
        fs,
        (4, 8),
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )


def eeg_alpha_power(
    x, fs, axis=-1, nperseg=None, scaling='density', return_db=False, **scipy_welch_kwargs
):
    """Alpha 波段功率 (8-12 Hz)

    计算 EEG 信号的 Alpha 波段 (8-12 Hz) 功率。

    Parameters
    ----------
    - `x` : array_like, 1-D or 2-D
        - 待处理的原始 EEG 信号, 可以是一维或二维数组, 类型为 `array_like`。
    - `fs` : float
        - 输入信号的采样频率, 单位为 Hz, 类型为 `float`。
    - `axis` : int, optional
        - 可选参数, 计算取样信号频谱的轴。默认值为 -1, 即最后一个轴, 类型为 `int`。
    - `nperseg` : int, optional
        - 每段的长度, 默认 2 * fs
    - `scaling` : str, optional
        - 'density': 返回功率谱密度 (功率/Hz)
        - 'spectrum': 返回功率谱 (功率)
    - `return_db` : bool, optional
        - 是否将功率谱转换为分贝 (dB) 形式

    Returns
    ----------
    - `band_power` : `float`
        - Alpha 波段功率。
    """
    return bandpower(
        x,
        fs,
        (8, 12),
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )


def eeg_sigma_power(
    x, fs, axis=-1, nperseg=None, scaling='density', return_db=False, **scipy_welch_kwargs
):
    """Sigma 波段功率 (12-15 Hz)

    计算 EEG 信号的 Sigma 波段 (12-15 Hz) 功率。

    Parameters
    ----------
    - `x` : array_like, 1-D or 2-D
        - 待处理的原始 EEG 信号, 可以是一维或二维数组, 类型为 `array_like`。
    - `fs` : float
        - 输入信号的采样频率, 单位为 Hz, 类型为 `float`。
    - `axis` : int, optional
        - 可选参数, 计算取样信号频谱的轴。默认值为 -1, 即最后一个轴, 类型为 `int`。
    - `nperseg` : int, optional
        - 每段的长度, 默认 2 * fs
    - `scaling` : str, optional
        - 'density': 返回功率谱密度 (功率/Hz)
        - 'spectrum': 返回功率谱 (功率)
    - `return_db` : bool, optional
        - 是否将功率谱转换为分贝 (dB) 形式

    Returns
    ----------
    - `band_power` : `float`
        - Sigma 波段功率。
    """
    return bandpower(
        x,
        fs,
        (12, 15),
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )


def eeg_beta_power(
    x, fs, axis=-1, nperseg=None, scaling='density', return_db=False, **scipy_welch_kwargs
):
    """Beta 波段功率 (12-30 Hz)

    计算 EEG 信号的 Beta 波段 (12-30 Hz) 功率。

    Parameters
    ----------
    - `x` : array_like, 1-D or 2-D
        - 待处理的原始 EEG 信号, 可以是一维或二维数组, 类型为 `array_like`。
    - `fs` : float
        - 输入信号的采样频率, 单位为 Hz, 类型为 `float`。
    - `axis` : int, optional
        - 可选参数, 计算取样信号频谱的轴。默认值为 -1, 即最后一个轴, 类型为 `int`。
        - `nperseg` : int, optional
        - 每段的长度, 默认 2 * fs
    - `scaling` : str, optional
        - 'density': 返回功率谱密度 (功率/Hz)
        - 'spectrum': 返回功率谱 (功率)
    - `return_db` : bool, optional
        - 是否将功率谱转换为分贝 (dB) 形式

    Returns
    ----------
    - `band_power` : `float`
        - Beta 波段功率。
    """
    return bandpower(
        x,
        fs,
        (12, 30),
        axis=axis,
        nperseg=nperseg,
        scaling=scaling,
        return_db=return_db,
        **scipy_welch_kwargs,
    )
