"""最大李雅普诺夫指数 (Largest Lyapunov Exponent)

使用 Rosenstein 算法计算最大李雅普诺夫指数。
李雅普诺夫指数衡量相空间中相邻轨道的指数分离速率：
- λ > 0：混沌行为
- λ ≈ 0：周期或准周期
- λ < 0：收敛到稳定不动点
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

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


# =============================================================================
# 最大李雅普诺夫指数 (Largest Lyapunov Exponent)
# =============================================================================


def _lyapunov_rosenstein_1d(
    x: np.ndarray, m: int, tau: int, min_tsep: int, max_iter: int, fs: float
) -> float:
    """Rosenstein 算法计算最大李雅普诺夫指数

    参考: Rosenstein et al. (1993) "A practical method for calculating largest
    Lyapunov exponents from small data sets"
    """
    N = len(x)
    if N < 50:
        return np.nan

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

    n_vectors = embedded.shape[0]
    if n_vectors < 2 * min_tsep + 1:
        return np.nan

    # 对每个点找到最近邻 (排除时间上太近的点)
    divergence = np.zeros(max_iter)
    count = np.zeros(max_iter)

    for i in range(n_vectors - max_iter):
        # 计算到所有其他点的距离
        distances = np.linalg.norm(embedded - embedded[i], axis=1)

        # 显式排除自身点（增强鲁棒性）
        distances[i] = np.inf

        # 排除时间上太近的点 (Theiler window)
        window_start = max(0, i - min_tsep)
        window_end = min(n_vectors, i + min_tsep + 1)
        distances[window_start:window_end] = np.inf

        # 找最近邻
        j = np.argmin(distances)
        if distances[j] == np.inf or distances[j] < 1e-10:
            continue

        # 跟踪轨道分离
        max_k = min(max_iter, n_vectors - i, n_vectors - j)
        for k in range(max_k):
            dist_k = np.linalg.norm(embedded[i + k] - embedded[j + k])
            if dist_k > 0:
                divergence[k] += np.log(dist_k)
                count[k] += 1

    # 计算平均对数分离
    valid_mask = count > 0
    if np.sum(valid_mask) < 5:
        return np.nan

    divergence[valid_mask] /= count[valid_mask]

    # 线性拟合斜率即为最大 Lyapunov 指数
    time_axis = np.arange(max_iter)[valid_mask] / fs
    div_valid = divergence[valid_mask]

    # 只取前半部分进行拟合（避免饱和区）
    n_fit = max(5, len(time_axis) // 2)
    try:
        coeffs = np.polyfit(time_axis[:n_fit], div_valid[:n_fit], 1)
        return coeffs[0]  # 斜率
    except (np.linalg.LinAlgError, ValueError):
        return np.nan


def largest_lyapunov_exponent(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    min_tsep: int | None = None,
    max_iter: int | None = None,
    fs: float = 1.0,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的最大李雅普诺夫指数 (Largest Lyapunov Exponent)

    使用 Rosenstein 算法。李雅普诺夫指数衡量相空间中相邻轨道的指数分离速率：
    - λ > 0：混沌行为
    - λ ≈ 0：周期或准周期
    - λ < 0：收敛到稳定不动点

    复杂度: O(N² × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟，默认自动估计, by default None
    - `min_tsep`: `int`, 寻找最近邻时的最小时间分离 (Theiler window)，默认 tau, by default None
    - `max_iter`: `int`, 跟踪轨道分离的最大迭代次数，默认 N//10, by default None
    - `fs`: `float`, 采样率，用于将 Lyapunov 指数转换为 1/s 单位, by default 1.0
    - `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, max_iter: int, fs: float
    ) -> float:
        N = len(x)
        _tau = tau if tau is not None else _estimate_tau_autocorr(x)
        _min_tsep = min_tsep if min_tsep is not None else _tau
        _max_iter = max_iter if max_iter is not None else max(10, N // 10)
        return _lyapunov_rosenstein_1d(x, m, _tau, _min_tsep, _max_iter, fs)

    return adapt_1d_func(signal, _wrapper, axis, m, tau, min_tsep, max_iter, fs)
