"""模糊熵 (Fuzzy Entropy)

模糊熵使用模糊隶属函数替代样本熵中的硬阈值判断，
提供更平滑的相似性度量，对参数选择更加鲁棒。
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func
from ._helper import _create_embedding_centered


# =============================================================================
# 模糊隶属函数
# =============================================================================


def _fuzzy_membership(d: np.ndarray, r: float, n: float = 2.0) -> np.ndarray:
    """模糊隶属函数: exp(-(d/r)^n)

    Parameters
    ----------
    - `d`: `np.ndarray`, 距离
    - `r`: `float`, 模糊边界宽度
    - `n`: `float`, 模糊指数, by default 2.0

    Returns
    ----------
    - `np.ndarray`, 隶属度 ∈ [0, 1]
    """
    return np.exp(-np.power(d / r, n))


# =============================================================================
# 模糊熵 (Fuzzy Entropy)
# =============================================================================


def _fuzzy_entropy_1d(x: np.ndarray, m: int, r_coeff: float, n: float) -> float:
    """计算一维信号的模糊熵

    归一化因子：
    - 原实现：sum / (N_vec - 1)
    - 修正后：2 * sum_upper / (N_vec * (N_vec - 1))
    """
    N = len(x)
    if N < m + 2:
        return np.nan

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

    r = cast(float, r_coeff * std_x)

    xm = _create_embedding_centered(x, m)
    xm1 = _create_embedding_centered(x, m + 1)

    if xm.shape[0] < 2 or xm1.shape[0] < 2:
        return np.nan

    def _compute_phi(vecs: np.ndarray) -> float:
        """计算平均相似度 φ

        归一化因子：
        φ = (2 / (N_vec * (N_vec - 1))) * Σ_{i<j} sim(i,j)
        """
        N_vec = len(vecs)
        if N_vec < 2:
            return 0.0

        # 计算上三角相似度之和
        sum_sim_upper = 0.0
        for i in range(N_vec - 1):
            dists = np.max(np.abs(vecs[i + 1 :] - vecs[i]), axis=1)
            sum_sim_upper += np.sum(_fuzzy_membership(dists, r, n))

        # [FIX] 正确归一化：全矩阵配对数（不含对角线）= N_vec * (N_vec - 1)
        # 上三角求和 * 2 / 全配对数
        n_pairs = N_vec * (N_vec - 1)
        return (2 * sum_sim_upper) / n_pairs

    phi_m = _compute_phi(xm)
    phi_m1 = _compute_phi(xm1)

    if phi_m == 0:
        return np.nan
    if phi_m1 == 0:
        return np.inf

    return np.log(phi_m) - np.log(phi_m1)


def fuzzy_entropy(
    signal: np.ndarray, m: int = 2, r: float = 0.2, n: float = 2.0, axis: int = -1
) -> np.ndarray:
    """计算信号的模糊熵 (Fuzzy Entropy)

    模糊熵使用模糊隶属函数替代样本熵中的硬阈值判断，
    提供更平滑的相似性度量，对参数选择更加鲁棒。

    复杂度: O(N² × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `r`: `float`, 模糊边界宽度系数（相对于标准差）, by default 0.2
    - `n`: `float`, 模糊指数（隶属函数的陡峭程度）, by default 2.0
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 模糊熵值
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _fuzzy_entropy_1d, axis, m, r, n)
