"""分形维度特征提取算法

特征列表：
- Higuchi 分形维度 (Higuchi Fractal Dimension)
- Katz 分形维度 (Katz Fractal Dimension)
- Petrosian 分形维度 (Petrosian Fractal Dimension)
- 相关维数 (Correlation Dimension)
"""

import numpy as np
from typing import Tuple, cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func
from ._helper import _embed_time_delay, _estimate_tau_autocorr


# =============================================================================
# Higuchi 分形维度
# =============================================================================


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.
                # This shouldn't happen if N >= 2*kmax, but handle anyway.
                L_mk = 0
            else:
                # Normalize L_m(k) as per Higuchi's definition
                # L_mk_normalized = L_mk * (N - 1) / (num_segments_in_subseq * k)
                # This is effectively averaging the segment lengths by N-1 / (actual_segments * k)
                L_mk = L_mk * (N - 1) / (num_segments_in_subseq * k)

            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]  # 斜率的负值即为分形维度


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`, 最大时间延迟/聚合级别, by default 10
    - `axis`: `int`, 计算轴, by default -1

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

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


# =============================================================================
# Katz 分形维度
# =============================================================================


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


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`, 计算轴, by default -1

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

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

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


# =============================================================================
# Petrosian 分形维度
# =============================================================================


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


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

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

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `axis`: `int`, 计算轴, by default -1

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

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


# =============================================================================
# 相关维数 (Correlation Dimension)
# =============================================================================


def _correlation_dimension_1d(
    x: np.ndarray,
    m: int,
    tau: int,
    min_tsep: int,  # Theiler 窗口参数
    r_range: Tuple[float, float],
    n_r: int,
) -> float:
    """Grassberger-Procaccia 算法计算相关维数

    增加 Theiler 窗口以消除时间自相关偏差
    """
    N = len(x)
    if N < 100:
        return np.nan

    # 相空间重构
    try:
        embedded = _embed_time_delay(x, m, tau)
    except ValueError:
        return np.nan

    n_vectors = embedded.shape[0]
    if n_vectors < 50:
        return np.nan

    # 估计信号范围以设置 r 的范围
    std_x = np.std(x)
    if std_x < 1e-10:
        return np.nan

    r_min = r_range[0] * std_x
    r_max = r_range[1] * std_x
    r_values = np.logspace(np.log10(r_min), np.log10(r_max), n_r)

    correlation_sum = np.zeros(n_r)

    if n_vectors <= 2000:
        # 全量计算距离矩阵
        dists = np.zeros((n_vectors, n_vectors))
        for i in range(n_vectors):
            dists[i] = np.linalg.norm(embedded - embedded[i], axis=1)

        # 使用 Theiler 窗口：排除 |i-j| <= min_tsep 的点对
        # triu_indices 的 k 参数控制对角线偏移
        valid_indices = np.triu_indices(n_vectors, k=max(1, min_tsep + 1))
        valid_dists = dists[valid_indices]

    else:
        # 大数据集使用采样
        n_sample = min(1000, n_vectors)
        sample_idx = np.random.choice(n_vectors, n_sample, replace=False)
        valid_dists_list = []

        for idx in sample_idx:
            d = np.linalg.norm(embedded - embedded[idx], axis=1)
            # 排除 Theiler 窗口内的点
            mask = np.abs(np.arange(n_vectors) - idx) > min_tsep
            valid_dists_list.append(d[mask])

        valid_dists = np.concatenate(valid_dists_list)

    if len(valid_dists) == 0:
        return np.nan

    n_pairs = len(valid_dists)

    for idx, r in enumerate(r_values):
        correlation_sum[idx] = np.sum(valid_dists < r) / n_pairs

    # 在双对数坐标下拟合斜率
    valid_mask = (correlation_sum > 0.01) & (correlation_sum < 0.99)
    if np.sum(valid_mask) < 3:
        return np.nan

    log_r = np.log(r_values[valid_mask])
    log_C = np.log(correlation_sum[valid_mask])

    try:
        coeffs = np.polyfit(log_r, log_C, 1)
        return coeffs[0]
    except (np.linalg.LinAlgError, ValueError):
        return np.nan


def correlation_dimension(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    min_tsep: int | None = None,
    r_range: Tuple[float, float] = (0.01, 1.0),
    n_r: int = 20,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的相关维数 (Correlation Dimension)

    使用 Grassberger-Procaccia 算法。增加了 min_tsep 参数 (Theiler window) 以消除时间自相关偏差。
    对于过采样信号，相邻点在相空间中距离很近，会导致 D₂ 被低估。

    复杂度: O(N² × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟，默认自动估计, by default None
    - `min_tsep`: `int`, Theiler 窗口大小，排除 |i-j| <= min_tsep 的点对，默认为 tau, by default None
    - `r_range`: `Tuple[float, float]`, 相对于信号标准差的 r 值范围, by default (0.01, 1.0)
    - `n_r`: `int`, r 值的采样点数, by default 20
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 相关维数
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))

    def _wrapper(
        x: np.ndarray, m: int, tau: int, min_tsep: int, r_range: Tuple, n_r: int
    ) -> float:
        _tau = tau if tau is not None else _estimate_tau_autocorr(x)
        _min_tsep = min_tsep if min_tsep is not None else max(1, _tau)
        return _correlation_dimension_1d(x, m, _tau, _min_tsep, r_range, n_r)

    return adapt_1d_func(signal, _wrapper, axis, m, tau, min_tsep, r_range, n_r)
