"""非线性特征提取算法。

提供各种熵度量和分形维度计算方法，用于捕捉脑电信号的非线性动态特性:

- 样本熵 (Sample Entropy)
- 近似熵 (Approximate Entropy)
- 排列熵 (Permutation Entropy)
- 多尺度熵 (Multiscale Entropy)
- Higuchi分形维度
- Katz分形维度
- Petrosian分形维度
- Hurst指数

#### References
----------
- [1] Richman, J. S., & Moorman, J. R. (2000). Physiological time-series analysis using
    approximate entropy and sample entropy. American Journal of Physiology-Heart and
    Circulatory Physiology, 278(6), H2039-H2049.
- [2] Bandt, C., & Pompe, B. (2002). Permutation entropy: a natural complexity measure
    for time series. Physical review letters, 88(17), 174102.
- [3] Costa, M., Goldberger, A. L., & Peng, C. K. (2005). Multiscale entropy analysis of
    biological signals. Physical review E, 71(2), 021906.
- [4] Higuchi, T. (1988). Approach to an irregular time series on the basis of the fractal
    theory. Physica D: Nonlinear Phenomena, 31(2), 277-283.
"""

import math

import numpy as np

# from scipy import stats
from zyplib.utils.ensure import ensure_ndarray

from .utils import adapt_1d_func


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_sampen(series: np.ndarray, dim: int) -> np.ndarray:
    n_pts = len(series)
    if n_pts < dim:
        return np.array([])  # Not enough points

    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 _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_sampen(x, m)
    xm1 = __create_embedding_sampen(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`, optional
        - 嵌入维度，默认为2
    - `r` : `float`, optional
        - 相似性阈值系数，实际阈值为 r * std(signal_segment)，默认为0.2
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

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

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


def __phi(x: np.ndarray, dim: int, r_thresh: float) -> float:
    # 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 = r_coeff * std_x
    if r_thresh < 0:
        r_thresh = 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`, optional
        - 嵌入维度，默认为2
    - `r` : `float`, optional
        - 相似性阈值系数，实际阈值为 r * std(signal_segment)，默认为0.2
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

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

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


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 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`, optional
        - 排列阶数，默认为3
    - `delay` : `int`, optional
        - 时间延迟，默认为1
    - `normalize` : `bool`, optional
        - 是否对排列熵进行归一化，默认为False
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

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

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


def higuchi_fractal_dimension(
    signal: np.ndarray, kmax: int = 10, axis: int = -1
) -> np.ndarray:
    """计算信号的Higuchi分形维度

    Higuchi分形维度是一种测量时间序列复杂度的方法，适用于非平稳信号。 O(N * kmax)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `kmax` : `int`, optional
        - 最大时间延迟/聚合级别，默认为10
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

    Returns
    -------
    - `np.ndarray`
        - Higuchi分形维度
    """
    signal = ensure_ndarray(signal)

    def _hfd_1d(x: np.ndarray, kmax: int) -> float:
        """计算一维信号的Higuchi分形维度"""
        N = len(x)
        if N < 2 * kmax or N < 20:  # Ensure N is large enough for meaningful calculation
            return np.nan  # Or some other indicator for short series

        lk = np.zeros(kmax)
        k_values = np.arange(1, kmax + 1)

        for k_idx, k in enumerate(k_values):
            lm_sum = 0.0
            # 对每个起始点 m 从 1 到 k (或 0 到 k-1)
            for m_offset in range(k):
                # 构建子序列 X_k^m
                # Indices: m, m+k, m+2k, ... up to N-1
                # In 0-based indexing: m_offset, m_offset+k, ...
                subseq = x[m_offset::k]

                if len(subseq) < 2:  # Need at least two points to calculate diff
                    # This case should ideally not happen if N is large enough relative to kmax
                    # but as a safeguard:
                    lm_sum += 0  # Or handle as error / skip this m_offset
                    continue

                # 计算子序列的长度 L_m(k)
                L_mk = np.sum(np.abs(np.diff(subseq)))

                # 归一化因子 (N-1) / ( floor((N-m)/k) * k )
                # N_sub = len(subseq) # Number of points in this specific sub-sequence
                # Normalization factor: (N-1) / ( (N_sub-1) * k ) if N_sub > 1
                # Higuchi's paper: L_m(k) = { (sum_{i=1}^{floor((N-m)/k)} |X(m+ik) - X(m+(i-1)k)|) * (N-1) } / { floor((N-m)/k) * k }

                # Number of segments in the sub-sequence
                num_segments_in_subseq = len(subseq) - 1
                if num_segments_in_subseq == 0:  # only one point in subseq
                    # This means (N-m_offset)/k was < 1, so subseq had 1 element.
                    # L_mk would be 0 (as diff is empty).
                    # This case should be avoided by ensuring N is large enough.
                    # If it happens, this m contributes 0 to lm_sum.
                    pass  # L_mk is already 0 from sum(abs(diff(single_point_array)))
                else:
                    if num_segments_in_subseq > 0:
                        L_mk = (L_mk * (N - 1)) / (num_segments_in_subseq * k)
                    else:  # subseq has 0 or 1 point
                        L_mk = 0.0

                lm_sum += L_mk

            lk[k_idx] = lm_sum / k  # Average L_m(k) over m

        # 在双对数坐标下拟合直线
        valid_mask = lk > 0
        if np.sum(valid_mask) < 2:  # Need at least 2 points for polyfit
            return np.nan

        log_lk = np.log(lk[valid_mask])
        log_k_values = np.log(k_values[valid_mask])

        # 线性回归
        try:
            polyfit_coeffs = np.polyfit(log_k_values, log_lk, 1)
        except (np.linalg.LinAlgError, ValueError):  # Handle cases where polyfit fails
            return np.nan

        return -polyfit_coeffs[0]  # 斜率的负值即为分形维度

    # 使用通用函数处理多维数组
    return adapt_1d_func(signal, _hfd_1d, axis, kmax)


def katz_fractal_dimension(signal: np.ndarray, axis: int = -1) -> np.ndarray:
    """计算信号的Katz分形维度

    Katz分形维度是另一种测量时间序列复杂度的方法。
    本实现使用n=N-1（段数）而非N（点数）计算，以符合Katz (1988)原始论文。 O(N)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

    Returns
    -------
    - `np.ndarray`
        - Katz分形维度

    Notes
    -----
    使用公式 log10(n) / (log10(d/L) + log10(n))，其中n=N-1为段数，
    以符合Katz (1988)原始论文。一些实现可能使用N（点数）。
    """
    signal = ensure_ndarray(signal)

    def _kfd_1d(x: np.ndarray) -> float:
        """计算一维信号的Katz分形维度"""
        N = len(x)
        if N < 2:
            return np.nan  # Need at least 2 points

        # 计算曲线长度 L = sum of distances between successive points
        # For 1D series, this is sum of absolute differences
        L = np.sum(np.abs(np.diff(x)))
        if L == 0:  # Constant signal
            return 1.0  # Dimension of a point or straight line

        # 计算从第一个点到序列中其他所有点的最大距离 d (diameter)
        # d = max_i ||x_1 - x_i||
        d_max = np.max(np.abs(x - x[0]))

        if d_max == 0:
            # This implies all points are same as x[0], so L should also be 0.
            return 1.0  # Dimension of a point

        # Katz分形维度公式: D = log10(n) / (log10(d/L) + log10(n))
        # 根据Katz (1988)原始论文，n是步数(N-1)而不是点数(N)
        # 使用n_segments = N-1更符合原始定义

        # 计算步数
        n_segments = N - 1

        # 对于N=2的情况，n_segments=1，log10(1)=0
        # 如果d_max/L=1（对于两点的情况通常如此），则分母为0
        # 对于直线（两点），通常认为其分形维度为1.0
        if n_segments == 1:  # N=2的情况
            return 1.0

        log_n = np.log10(n_segments)

        # 计算分母
        denominator = np.log10(d_max / L) + log_n

        if denominator == 0:  # 避免除以零
            return np.nan

        return log_n / denominator

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


def petrosian_fractal_dimension(signal: np.ndarray, axis: int = -1) -> np.ndarray:
    """计算信号的Petrosian分形维度

    Petrosian分形维度是一种简化的分形维度计算方法。 O(N)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

    Returns
    -------
    - `np.ndarray`
        - Petrosian分形维度
    """
    signal = ensure_ndarray(signal)

    def _pfd_1d(x: np.ndarray) -> float:
        """计算一维信号的Petrosian分形维度"""
        N = len(x)
        if N < 2:  # Need at least 2 points for diff
            return np.nan

        # 计算信号的一阶差分
        diff_x = np.diff(x)
        if len(diff_x) < 2:  # Need at least 2 differences for sign_changes calculation
            return np.nan  # 序列过短无法计算

        # 计算符号变化的次数 N_delta
        # sign_changes = np.sum(np.sign(diff_x[:-1]) != np.sign(diff_x[1:]))
        # The original formula counts where product is negative (diff signs, excluding zeros)
        sign_changes = np.sum(diff_x[:-1] * diff_x[1:] < 0)

        # Petrosian分形维度公式: D = log10(N) / (log10(N) + log10(N / (N + C * N_delta)))
        # where C is often 0.4 as in the original code.
        log_N = np.log10(N)

        # N > 0 且 sign_changes >= 0，因此 denominator_log_arg 总是 > 0
        denominator_log_arg = N / (N + 0.4 * sign_changes)
        pfd = log_N / (log_N + np.log10(denominator_log_arg))
        return pfd

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


def _hurst_1d(x: np.ndarray, max_lag_arg: int) -> float:
    """计算一维信号的Hurst指数"""
    N = len(x)

    # Adjust max_lag for short series
    # Ensure max_lag is at least 2 and not too large for N
    if N < 20:  # If N is very small, Hurst exponent is unreliable
        return np.nan

    # Ensure max_lag is appropriate for N
    # Common rule: max_lag <= N/2
    current_max_lag = min(max_lag_arg, N // 2)
    current_max_lag = max(2, current_max_lag)  # lags start from 2

    if current_max_lag < 2:  # Not enough lags to compute
        return np.nan

    # lags for R/S calculation (n in many texts)
    lags = np.arange(2, current_max_lag + 1)  # Corrected to include max_lag
    if len(lags) < 2:  # Need at least two (lag, R/S) points for regression
        return np.nan

    # 计算每个滞后时间的重标范围 (R/S value)
    rs_values = np.zeros(len(lags))

    for i, lag_val in enumerate(lags):
        # 将序列分成多个不重叠的子序列，每个长度为 lag_val
        n_blocks = N // lag_val
        if n_blocks < 1:  # Should not happen due to current_max_lag <= N//2
            rs_values[i] = np.nan
            continue

        # Reshape into blocks
        sub_series_matrix = x[: n_blocks * lag_val].reshape(n_blocks, lag_val)

        # 1. 计算每个子序列的均值调整序列 (Y_t = X_t - mean(X))
        mean_adjusted_blocks = sub_series_matrix - np.mean(
            sub_series_matrix, axis=1, keepdims=True
        )

        # 2. 计算累积离差 Z_k = sum_{t=1 to k} Y_t
        cumulative_deviates = np.cumsum(mean_adjusted_blocks, axis=1)

        # 3. 计算每个子序列的范围 R = max(Z) - min(Z)
        block_ranges = np.max(cumulative_deviates, axis=1) - np.min(
            cumulative_deviates, axis=1
        )

        # 4. 计算每个子序列的标准差 S
        # ddof=0 使用总体标准差计算 (基于N个元素)
        # ddof=1 使用样本标准差计算 (基于N-1个元素)
        # 这里显式指定ddof=0，与之前的隐式行为保持一致
        block_stds = np.std(sub_series_matrix, axis=1, ddof=0)

        # 5. 计算每个子序列的 R/S 值
        # Avoid division by zero if a block has std=0 (constant block)
        # If std is 0, range is also 0. R/S = 0/0 = nan.
        # Or, if std is 0, R/S is often taken as 0 or excluded.
        rs_block_ratios = np.zeros(n_blocks)
        for j in range(n_blocks):
            if block_stds[j] < 1e-10:  # Effectively zero std
                if block_ranges[j] < 1e-10:  # Range is also zero
                    rs_block_ratios[j] = 0.0  # R/S = 0 for constant block
                else:  # Std is zero but range is not (should not happen for mean-adjusted)
                    rs_block_ratios[j] = np.nan  # Undefined
            else:
                rs_block_ratios[j] = block_ranges[j] / block_stds[j]

        # 6. 平均 R/S 值
        valid_rs_ratios = rs_block_ratios[~np.isnan(rs_block_ratios)]
        if len(valid_rs_ratios) > 0:
            rs_values[i] = np.mean(valid_rs_ratios)
        else:
            rs_values[i] = np.nan

    # 在双对数坐标下拟合直线: log(R/S) vs log(lag)
    valid_mask = ~np.isnan(rs_values) & (rs_values > 1e-10)  # Ensure R/S > 0 for log

    if np.sum(valid_mask) < 2:  # Need at least 2 points for polyfit
        return np.nan

    log_rs = np.log(rs_values[valid_mask])
    log_lags = np.log(lags[valid_mask])

    try:
        polyfit_coeffs = np.polyfit(log_lags, log_rs, 1)
    except (np.linalg.LinAlgError, ValueError):
        return np.nan

    return polyfit_coeffs[0]  # 斜率即为Hurst指数


def hurst_exponent(signal: np.ndarray, max_lag: int = 20, axis: int = -1) -> np.ndarray:
    """计算信号的Hurst指数

    Hurst指数用于评估时间序列的长期记忆性和分形特性。计算比较慢。 O(N·max_lag)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `max_lag` : `int`, optional
        - 最大滞后时间 (inclusive)，默认为20
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度

    Returns
    -------
    - `np.ndarray`
        - Hurst指数
    """
    signal = ensure_ndarray(signal)

    # 使用通用函数处理多维数组
    return adapt_1d_func(signal, _hurst_1d, axis, max_lag)
