"""谱域熵度量算法

特征列表：
- 谱熵 (Spectral Entropy)
- 小波熵 (Wavelet Entropy)
"""

import warnings
from typing import cast

import numpy as np

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func


# =============================================================================
# 谱熵 (Spectral Entropy)
# =============================================================================


def _spectral_entropy_1d(
    x: np.ndarray, fs: float, method: str, nperseg: int, normalize: bool
) -> float:
    """计算一维信号的谱熵"""
    from scipy import signal as sig

    N = len(x)
    if N < 10:
        return np.nan

    # 计算功率谱密度
    if method == 'welch':
        _nperseg = nperseg if nperseg is not None else min(256, N)
        f, psd = sig.welch(x, fs=fs, nperseg=_nperseg)
    elif method == 'periodogram':
        f, psd = sig.periodogram(x, fs=fs)
    else:
        _nperseg = nperseg if nperseg is not None else min(256, N)
        f, psd = sig.welch(x, fs=fs, nperseg=_nperseg)

    psd_sum = np.sum(psd)
    if psd_sum == 0:
        return 0.0

    psd_norm = psd / psd_sum
    psd_norm = psd_norm[psd_norm > 0]

    entropy = -np.sum(psd_norm * np.log2(psd_norm))

    if normalize:
        max_entropy = np.log2(len(psd_norm))
        if max_entropy > 0:
            entropy /= max_entropy

    return float(entropy)


def spectral_entropy(
    signal: np.ndarray,
    fs: float = 1.0,
    method: str = 'welch',
    nperseg: int | None = None,
    normalize: bool = True,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的谱熵 (Spectral Entropy)

    谱熵度量功率谱的平坦程度，反映信号的频率复杂度。

    复杂度: O(N log N)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `fs`: `float`, 采样率, by default 1.0
    - `method`: `str`, 功率谱估计方法，'welch' 或 'periodogram', by default 'welch'
    - `nperseg`: `int`, Welch 方法的段长度, by default None
    - `normalize`: `bool`, 是否归一化到 [0, 1], by default True
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 谱熵值
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(
        signal, _spectral_entropy_1d, axis, fs, method, nperseg, normalize
    )


# =============================================================================
# 小波熵 (Wavelet Entropy)
# =============================================================================


def _wavelet_entropy_1d(
    x: np.ndarray, wavelet: str, level: int, normalize: bool
) -> float:
    """计算一维信号的小波熵"""
    try:
        import pywt
    except ImportError:
        warnings.warn('小波熵需要 pywt 库：pip install PyWavelets')
        return np.nan

    N = len(x)
    if N < 2**level:
        return np.nan

    coeffs = pywt.wavedec(x, wavelet, level=level)

    # 计算各尺度能量
    energies = np.array([np.sum(c**2) for c in coeffs])
    total_energy = np.sum(energies)

    if total_energy == 0:
        return 0.0

    probs = energies / total_energy
    probs = probs[probs > 0]

    entropy = -np.sum(probs * np.log2(probs))

    if normalize:
        max_entropy = np.log2(len(coeffs))
        if max_entropy > 0:
            entropy /= max_entropy

    return entropy


def wavelet_entropy(
    signal: np.ndarray,
    wavelet: str = 'db4',
    level: int = 5,
    normalize: bool = False,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的小波熵 (Wavelet Entropy)

    小波熵基于信号的小波分解，度量能量在不同尺度上的分布均匀程度。

    复杂度: O(N × level)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `wavelet`: `str`, 小波类型, by default 'db4'
    - `level`: `int`, 分解层数, by default 5
    - `normalize`: `bool`, 是否归一化到 [0, 1], by default False
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 小波熵值
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _wavelet_entropy_1d, axis, wavelet, level, normalize)
