"""多尺度熵 (Multiscale Entropy, MSE)

通过在不同时间尺度上计算样本熵，揭示信号的多尺度复杂度结构。
原始提出用于分析心率变异性 (Costa et al., 2002)。
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func
from ._helper import _coarse_grain, _sample_entropy_core


# =============================================================================
# 多尺度熵 (Multiscale Entropy, MSE)
# =============================================================================


def _mse_1d(x: np.ndarray, m: int, r_coeff: float, max_scale: int) -> np.ndarray:
    """计算一维信号的多尺度熵

    注：阈值基于原始信号的标准差，符合 Costa et al. (2002) 原始定义
    """
    N = len(x)
    if N < 50:
        return np.full(max_scale, np.nan)

    std_x = np.std(x)
    if std_x == 0:
        return np.zeros(max_scale)

    # 阈值基于原始信号（这是正确的）
    r = cast(float, r_coeff * std_x)
    mse_values = np.zeros(max_scale)

    for scale in range(1, max_scale + 1):
        coarse = _coarse_grain(x, scale)
        if len(coarse) < m + 2:
            mse_values[scale - 1] = np.nan
        else:
            mse_values[scale - 1] = _sample_entropy_core(coarse, m, r)

    return mse_values


def multiscale_entropy(
    signal: np.ndarray, m: int = 2, r: float = 0.15, max_scale: int = 20, axis: int = -1
) -> np.ndarray:
    """计算信号的多尺度熵 (Multiscale Entropy, MSE)

    多尺度熵通过在不同时间尺度上计算样本熵，揭示信号的多尺度复杂度结构。
    原始提出用于分析心率变异性 (Costa et al., 2002)。

    复杂度: O(max_scale × N² × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `r`: `float`, 相似性阈值系数（相对于原始信号标准差）, by default 0.15
    - `max_scale`: `int`, 最大尺度, by default 20
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 各尺度的样本熵值，形状为 (..., max_scale)
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _mse_1d, axis, m, r, max_scale)


def multiscale_entropy_mean(
    signal: np.ndarray, m: int = 2, r: float = 0.15, max_scale: int = 20, axis: int = -1
) -> np.ndarray:
    """计算多尺度熵的均值作为单一复杂度指标

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `r`: `float`, 相似性阈值系数, by default 0.15
    - `max_scale`: `int`, 最大尺度, by default 20
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 多尺度熵均值
    """
    mse = multiscale_entropy(signal, m, r, max_scale, axis)
    return np.nanmean(mse, axis=-1)
