"""分散熵 (Dispersion Entropy)

分散熵使用正态 CDF 进行信号映射，比排列熵更好地利用幅度信息。
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func


# =============================================================================
# 分散熵 (Dispersion Entropy)
# =============================================================================


def _dispersion_entropy_1d(
    x: np.ndarray, m: int, c: int, delay: int, normalize: bool
) -> float:
    """计算一维信号的分散熵"""
    from scipy.stats import norm

    N = len(x)
    # 检查长度是否足够
    if N < (m - 1) * delay + 1:
        return np.nan

    sigma = np.std(x)
    if sigma == 0:
        return 0.0

    # NCDF 映射到 [0, 1]
    x_normalized = (x - np.mean(x)) / sigma
    x_mapped = norm.cdf(x_normalized)

    # 映射到 {1, 2, ..., c}
    x_classes = np.clip(np.ceil(c * x_mapped), 1, c).astype(int)

    # 模式提取
    n_patterns = N - (m - 1) * delay

    # 使用索引矩阵向量化
    idx = np.arange(n_patterns)[:, None] + np.arange(m)[None, :] * delay
    patterns_mat = x_classes[idx]

    # 编码为整数（base-c 编码）
    powers = np.power(c, np.arange(m)[::-1])
    codes = np.dot(patterns_mat - 1, powers)

    _, counts = np.unique(codes, return_counts=True)
    probs = counts / n_patterns

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

    if normalize:
        max_entropy = np.log2(c**m)
        if max_entropy > 0:
            entropy /= max_entropy

    return entropy


def dispersion_entropy(
    signal: np.ndarray,
    m: int = 2,
    c: int = 6,
    delay: int = 1,
    normalize: bool = True,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的分散熵 (Dispersion Entropy)

    分散熵使用正态 CDF 进行信号映射，比排列熵更好地利用幅度信息。

    复杂度: O(N × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `c`: `int`, 类别数, by default 6
    - `delay`: `int`, 时间延迟, by default 1
    - `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, _dispersion_entropy_1d, axis, m, c, delay, normalize)
