import typing
from dataclasses import dataclass
from typing import Dict, Optional, Union

import numpy as np
import pywt
from typing_extensions import deprecated


@dataclass
class BandRange:
    """频带范围定义"""

    min_freq: float
    max_freq: float


CLASSIC_BANDS = {
    'delta': BandRange(0.5, 4.0),
    'theta': BandRange(4.0, 8.0),
    'alpha': BandRange(8.0, 13.0),
    'beta': BandRange(13.0, 30.0),
    'gamma': BandRange(30.0, 100.0),
}


def _validate_wavelet_inputs(
    signal: np.ndarray, wavelet: str, level: int, axis: int
) -> None:
    """验证小波变换的输入参数

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `wavelet` : `str`
        - 小波基函数
    - `level` : `int`
        - 分解层数
    - `axis` : `int`
        - 计算轴

    Raises
    ------
    - `ValueError`
        - 如果输入参数无效
    """
    # 检查小波基是否有效
    if not pywt.Wavelet(wavelet):  # pyright: ignore[reportAttributeAccessIssue]
        raise ValueError(f'无效的小波基函数: {wavelet}')

    # 检查信号长度是否满足分解要求
    min_length = 2**level
    if signal.shape[axis] < min_length:
        raise ValueError(
            f'信号长度({signal.shape[axis]})必须大于等于2^level({min_length})'
        )

    # 检查分解层数
    max_level = pywt.dwt_max_level(
        signal.shape[axis],
        pywt.Wavelet(wavelet).dec_len,  # pyright: ignore[reportAttributeAccessIssue]
    )
    if level > max_level:
        raise ValueError(f'分解层数({level})超过了最大可能层数({max_level})')


def wavelet_features(
    signal: np.ndarray,
    wavelet: str = 'db4',
    level: int = 4,
    axis: int = -1,
    features: Optional[typing.Sequence] = None,
    epsilon: float = 1e-10,
) -> Dict[str, np.ndarray]:
    """计算小波特征

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `wavelet` : `str`, optional
        - 小波基函数，默认为'db4'
    - `level` : `int`, optional
        - 分解层数，默认为4
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度
    - `features` : `list`, optional
        - 需要计算的特征列表，可选['mean', 'std', 'energy', 'entropy']
        - 默认为全部特征
    - `epsilon` : `float`, optional
        - 计算熵时的平滑参数，默认为1e-10

    Returns
    -------
    - `Dict[str, np.ndarray]`
        - 包含各层小波系数的统计特征
    """
    # 验证输入
    _validate_wavelet_inputs(signal, wavelet, level, axis)

    # 默认计算所有特征
    if features is None:
        features = ['mean', 'std', 'energy', 'entropy']

    # 小波分解
    coeffs = pywt.wavedec(signal, wavelet, level=level, axis=axis)

    result = {}

    # 特征计算函数映射
    feature_funcs = {
        'mean': lambda x: np.mean(np.abs(x), axis=axis),
        'std': lambda x: np.std(x, axis=axis),
        'energy': lambda x: np.sum(x**2, axis=axis),
        'entropy': lambda x: -np.sum(
            np.where(x**2 > epsilon, x**2 * np.log(x**2), 0), axis=axis
        ),
    }

    # 对每层系数计算特征
    for i, coef in enumerate(coeffs):
        name = 'approx' if i == 0 else f'detail_{i}'

        for feat in features:
            if feat in feature_funcs:
                result[f'{name}_{feat}'] = feature_funcs[feat](coef)

    return result
