"""熵计算的辅助函数

提供熵算法所需的公共辅助函数：
- 时间序列嵌入
- 粗粒化处理
- 样本熵核心计算
"""

import numpy as np


# =============================================================================
# 时间序列嵌入
# =============================================================================


def _embed_seq(seq: np.ndarray, tau: int, de: int) -> np.ndarray:
    """将时间序列嵌入到高维空间

    Parameters
    ----------
    - `seq`: `np.ndarray`, 输入序列 (1D)
    - `tau`: `int`, 时间延迟
    - `de`: `int`, 嵌入维度

    Returns
    ----------
    - `np.ndarray`, 嵌入后的序列，形状为 (n_vectors, de)，其中 n_vectors = N - (de - 1) * tau
    """
    N = len(seq)
    n_vectors = N - (de - 1) * tau
    if n_vectors < 1:  # 修正后的错误检查
        raise ValueError('序列长度不足以进行指定维度和延迟的嵌入')

    Y = np.zeros((n_vectors, de))
    for i in range(de):
        Y[:, i] = seq[i * tau : i * tau + n_vectors]
    return Y


def _create_embedding(series: np.ndarray, dim: int) -> np.ndarray:
    """创建嵌入矩阵 (stride trick 实现)

    Parameters
    ----------
    - `series`: `np.ndarray`, 输入序列
    - `dim`: `int`, 嵌入维度

    Returns
    ----------
    - `np.ndarray`, 嵌入矩阵，形状为 (n_pts - dim + 1, dim)
    """
    n_pts = len(series)
    if n_pts < dim:
        return np.array([]).reshape(0, dim)
    shape = (n_pts - dim + 1, dim)
    strides = (series.strides[0], series.strides[0])
    return np.lib.stride_tricks.as_strided(series, shape=shape, strides=strides)


def _create_embedding_centered(series: np.ndarray, dim: int) -> np.ndarray:
    """创建去均值的嵌入向量

    Parameters
    ----------
    - `series`: `np.ndarray`, 输入序列
    - `dim`: `int`, 嵌入维度

    Returns
    ----------
    - `np.ndarray`, 去均值的嵌入矩阵
    """
    n_pts = len(series)
    if n_pts < dim:
        return np.array([]).reshape(0, dim)
    shape = (n_pts - dim + 1, dim)
    strides = (series.strides[0], series.strides[0])
    embedded = np.lib.stride_tricks.as_strided(
        series, shape=shape, strides=strides
    ).copy()
    # 去除每个向量的均值
    return embedded - embedded.mean(axis=1, keepdims=True)


# =============================================================================
# 粗粒化处理 (用于多尺度熵)
# =============================================================================


def _coarse_grain(x: np.ndarray, scale: int) -> np.ndarray:
    """粗粒化：将序列按 scale 分段平均

    Parameters
    ----------
    - `x`: `np.ndarray`, 输入序列
    - `scale`: `int`, 尺度因子

    Returns
    ----------
    - `np.ndarray`, 粗粒化后的序列，长度为 len(x) // scale
    """
    N = len(x)
    n_segments = N // scale
    if n_segments < 1:
        return np.array([])
    x_truncated = x[: n_segments * scale]
    return x_truncated.reshape(n_segments, scale).mean(axis=1)


# =============================================================================
# 样本熵核心计算
# =============================================================================


def _sample_entropy_core(x: np.ndarray, m: int, r: float) -> float:
    """样本熵核心计算，r 为绝对阈值

    用于多尺度熵等需要在粗粒化序列上计算样本熵的场景。

    Parameters
    ----------
    - `x`: `np.ndarray`, 输入序列
    - `m`: `int`, 嵌入维度
    - `r`: `float`, 绝对阈值

    Returns
    ----------
    - `float`, 样本熵值
    """
    N = len(x)
    if N < m + 2:
        return np.nan

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

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

    # 计算匹配数 B (m维) - 只计算上三角 (i < j)
    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)

    # 计算匹配数 A (m+1维)
    A = 0.0
    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)

    if B == 0:
        return np.nan
    if A == 0:
        return np.inf

    return -np.log(A / B)
