"""基础熵度量算法

特征列表：
- 样本熵 (Sample Entropy)
- 近似熵 (Approximate Entropy)
- 排列熵 (Permutation Entropy)
"""

import math

import numpy as np

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func
from ._helper import _embed_seq, _create_embedding


# =============================================================================
# 样本熵 (Sample Entropy)
# =============================================================================


def _sampen_1d(x: np.ndarray, m: int, r_coeff: float) -> float:
    """计算一维信号的样本熵"""
    N = len(x)
    if N == 0:
        return np.nan  # 空序列无法计算

    # 计算标准差和实际阈值
    std_x = np.std(x)  # ddof=0 by default
    if std_x == 0:  # 对于常数序列
        # 所有距离都为0，如果 r_thresh 也为0，则所有都匹配
        # A 和 B 会很大，但 A/B 接近1，log(A/B) 接近0
        # 这是期望的行为，表示确定性高，熵低
        # 如果 r_coeff > 0, r_thresh = 0. dist <= 0 意味着 dist=0.
        return 0.0

    r_thresh = r_coeff * std_x

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

    if xm.shape[0] == 0 or xm1.shape[0] == 0:
        return np.nan  # Not enough data points to form vectors

    # 计算匹配数 B (m维)

    B = 0.0
    for i in range(xm.shape[0] - 1):
        dists_b = np.max(np.abs(xm[i + 1 :] - xm[i]), axis=1)
        B += np.sum(dists_b <= r_thresh)

    A = 0.0
    # Iterate through xm1 vectors
    for i in range(xm1.shape[0] - 1):
        dists_a = np.max(np.abs(xm1[i + 1 :] - xm1[i]), axis=1)
        A += np.sum(dists_a <= r_thresh)

    # 计算样本熵
    if B == 0:
        return np.nan
    if A == 0:  # B > 0
        return np.inf  # Maximal irregularity or undefined

    return -np.log(A / B)


def sample_entropy(
    signal: np.ndarray, m: int = 2, r: float = 0.2, axis: int = -1
) -> np.ndarray:
    """计算信号的样本熵

    样本熵是一种测量时间序列复杂度的方法，值越低表示信号的规律性越强。O(N² * m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `r`: `float`, 相似性阈值系数，实际阈值为 r * std(signal_segment), by default 0.2
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 样本熵值
    """
    signal = np.asarray(ensure_ndarray(signal))

    # 使用通用函数处理多维数组
    return adapt_1d_func(signal, _sampen_1d, axis, m, r)


# =============================================================================
# 近似熵 (Approximate Entropy)
# =============================================================================


def _phi(x: np.ndarray, dim: int, r_thresh: float) -> float:
    """计算近似熵的 phi 函数"""
    # 1. 创建嵌入向量 (tau=1 for ApEn)
    try:
        # _embed_seq(seq, tau, de) returns (N - (de-1)*tau, de)
        # For ApEn, N_emb = N - dim + 1
        x_embedded = _embed_seq(x, 1, dim)
    except ValueError:
        return np.nan

    N_emb = x_embedded.shape[0]
    if N_emb == 0:
        return np.nan

    # 2. 计算 C_i^dim(r_thresh) 的对数和的平均值
    # C_i^dim(r) = (number of j such that d[X_dim(i), X_dim(j)] <= r_thresh) / N_emb
    log_C_sum = 0.0
    for i in range(N_emb):
        # 计算与 x_embedded[i] 的切比雪夫距离
        dists = np.max(np.abs(x_embedded - x_embedded[i]), axis=1)
        # 计算有多少个向量与 x_embedded[i] 相似 (包括自身)
        num_similar = np.sum(dists <= r_thresh)

        if num_similar == 0:
            # Avoid log(0).
            return np.nan

        C_i_val = num_similar / N_emb
        log_C_sum += np.log(C_i_val)

    return log_C_sum / N_emb


def _apen_1d(x: np.ndarray, m: int, r_coeff: float) -> float:
    """计算一维信号的近似熵"""
    N = len(x)
    if N == 0:
        return np.nan

    std_x = np.std(x)
    if std_x == 0:  # 对于常数序列，ApEn 应该是0
        return 0.0

    r_thresh = float(r_coeff * std_x)
    if r_thresh < 0:
        r_thresh = 0.0

    phi_m = _phi(x, m, r_thresh)
    phi_m_plus_1 = _phi(x, m + 1, r_thresh)

    if np.isnan(phi_m) or np.isnan(phi_m_plus_1):
        return np.nan

    # ApEn = phi^m(r) - phi^{m+1}(r)
    apen = phi_m - phi_m_plus_1
    return apen


def approximate_entropy(
    signal: np.ndarray, m: int = 2, r: float = 0.2, axis: int = -1
) -> np.ndarray:
    """计算信号的近似熵 (修正版)

    近似熵是样本熵的前身，计算速度更快但可能不如样本熵稳定。O(N² * m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 2
    - `r`: `float`, 相似性阈值系数，实际阈值为 r * std(signal_segment), by default 0.2
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 近似熵值
    """
    signal = np.asarray(ensure_ndarray(signal))

    # 使用通用函数处理多维数组
    return adapt_1d_func(signal, _apen_1d, axis, m, r)


# =============================================================================
# 排列熵 (Permutation Entropy)
# =============================================================================


def _perm_entropy_1d(x: np.ndarray, order: int, delay: int, normalize: bool) -> float:
    """计算一维信号的排列熵"""
    N = len(x)
    if N == 0:
        return np.nan

    try:
        embedded = _embed_seq(x, delay, order)
    except ValueError:  # Not enough points
        return np.nan

    if embedded.shape[0] == 0:  # No embedded vectors
        return np.nan

    # 获取每个嵌入向量的排列 (将排列模式映射到一个整数)
    # 使用 argsort 获得排序后的索引，然后基于这些索引计算一个唯一的整数值
    perms = np.apply_along_axis(np.argsort, 1, embedded)

    motifs = [tuple(row_perm) for row_perm in perms]

    _, counts = np.unique(motifs, return_counts=True, axis=0)

    # 计算概率
    probs = counts / len(embedded)

    # 计算熵 (log base 2 for bits, or log base e for nats)
    pe = -np.sum(probs * np.log2(probs))

    if normalize:
        # 避免除以0
        normalizer = math.log2(math.factorial(order))
        if normalizer > 0:  # Avoid division by zero if order=0 or 1
            pe /= normalizer

    return float(pe)


def permutation_entropy(
    signal: np.ndarray,
    order: int = 3,
    delay: int = 1,
    normalize: bool = False,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的排列熵

    排列熵基于信号中的序列模式，对噪声更加鲁棒。O(N * order * log(order))

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `order`: `int`, 排列阶数, by default 3
    - `delay`: `int`, 时间延迟, by default 1
    - `normalize`: `bool`, 是否对排列熵进行归一化, by default False
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 排列熵值
    """
    signal = np.asarray(ensure_ndarray(signal))

    # 使用通用函数处理多维数组
    return adapt_1d_func(signal, _perm_entropy_1d, axis, order, delay, normalize)
