"""递归定量分析 (Recurrence Quantification Analysis, RQA)

RQA 基于递归图 (Recurrence Plot) 提取定量指标，用于刻画动力系统的复杂度和结构。

特征列表：
- 递归率 (Recurrence Rate, RR)
- 确定性 (Determinism, DET)
- 平均对角线长度 (Average Diagonal Line Length, L)
- 最大对角线长度 (Longest Diagonal Line, Lmax)
- 熵 (Entropy of Diagonal Lines, ENTR)
- 层叠性 (Laminarity, LAM)
- 捕获时间 (Trapping Time, TT)
"""

import typing

import numpy as np
from typing import Dict, Union

from zyplib.utils.ensure import ensure_ndarray

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


# =============================================================================
# 递归矩阵计算
# =============================================================================


def _compute_recurrence_matrix(
    x: np.ndarray, m: int, tau: int, threshold: float
) -> np.ndarray:
    """计算递归矩阵"""
    embedded = _embed_time_delay(x, m, tau)
    n_vectors = embedded.shape[0]

    dist_matrix = np.zeros((n_vectors, n_vectors))
    for i in range(n_vectors):
        dist_matrix[i] = np.linalg.norm(embedded - embedded[i], axis=1)

    return dist_matrix <= threshold


def _get_run_lengths(binary_array: np.ndarray) -> np.ndarray:
    """高效获取二值数组中连续 1 的长度"""
    if len(binary_array) == 0:
        return np.array([], dtype=int)

    # 使用 diff 寻找边缘
    padded = np.concatenate(([0], binary_array.astype(int), [0]))
    diffs = np.diff(padded)
    starts = np.where(diffs == 1)[0]
    ends = np.where(diffs == -1)[0]
    return ends - starts


# =============================================================================
# 对角线和垂直线提取
# =============================================================================


def _extract_diagonal_lines(
    recurrence_matrix: np.ndarray, min_length: int = 2
) -> np.ndarray:
    """提取递归矩阵中的对角线长度（仅上三角，排除主对角线）"""
    n = recurrence_matrix.shape[0]
    diagonal_lengths = []

    for k in range(1, n):
        diag = np.diag(recurrence_matrix, k)
        lengths = _get_run_lengths(diag)
        diagonal_lengths.extend(lengths[lengths >= min_length])

    return np.array(diagonal_lengths) if diagonal_lengths else np.array([], dtype=int)


def _extract_vertical_lines(
    recurrence_matrix: np.ndarray, min_length: int = 2
) -> np.ndarray:
    """提取递归矩阵中的垂直线长度（全矩阵）"""
    n = recurrence_matrix.shape[0]
    vertical_lengths = []

    for j in range(n):
        col = recurrence_matrix[:, j]
        lengths = _get_run_lengths(col)
        vertical_lengths.extend(lengths[lengths >= min_length])

    return np.array(vertical_lengths) if vertical_lengths else np.array([], dtype=int)


# =============================================================================
# RQA 核心计算
# =============================================================================


def _rqa_1d(
    x: np.ndarray,
    m: int,
    tau: int,
    threshold: float,
    min_diag_length: int,
    min_vert_length: int,
) -> Dict[str, float]:
    """计算一维信号的 RQA 指标

    修正 LAM 计算：垂直线统计基于全矩阵，分母也使用全矩阵递归点数
    """
    N = len(x)
    nan_result = {
        'RR': np.nan,
        'DET': np.nan,
        'L': np.nan,
        'Lmax': np.nan,
        'ENTR': np.nan,
        'LAM': np.nan,
        'TT': np.nan,
    }

    if N < 50:
        return nan_result

    try:
        R = _compute_recurrence_matrix(x, m, tau, threshold)
    except ValueError:
        return nan_result

    n = R.shape[0]

    # 全矩阵递归点数
    n_recurrent_total = np.sum(R)
    # 主对角线递归点数（通常全为 1）
    n_diag_main = np.sum(np.diag(R))
    # 非主对角线递归点数
    n_recurrent_off_diag = n_recurrent_total - n_diag_main

    # 1. 递归率 RR（全矩阵）
    RR = n_recurrent_total / (n * n)

    # 2. 对角线统计（上三角，排除主对角线）
    diagonal_lengths = _extract_diagonal_lines(R, min_diag_length)
    sum_diag_upper = np.sum(diagonal_lengths) if len(diagonal_lengths) > 0 else 0

    # [FIX] DET: 对角线结构点数 / 非主对角线递归点数
    # 由于只计算了上三角，需要乘 2（矩阵对称）
    if n_recurrent_off_diag > 0:
        DET = (2 * sum_diag_upper) / n_recurrent_off_diag
        DET = min(DET, 1.0)
    else:
        DET = 0.0

    # L, Lmax, ENTR
    if len(diagonal_lengths) > 0:
        L = np.mean(diagonal_lengths)
        Lmax = np.max(diagonal_lengths)

        # 对角线长度分布熵
        max_len = int(Lmax)
        hist, _ = np.histogram(diagonal_lengths, bins=range(min_diag_length, max_len + 2))
        hist = hist[hist > 0]
        if len(hist) > 0:
            probs = hist / np.sum(hist)
            ENTR = -np.sum(probs * np.log(probs))
        else:
            ENTR = 0.0
    else:
        L, Lmax, ENTR = 0.0, 0, 0.0

    # 3. 垂直线统计（全矩阵）
    vertical_lengths = _extract_vertical_lines(R, min_vert_length)
    sum_vert = np.sum(vertical_lengths) if len(vertical_lengths) > 0 else 0

    # [FIX] LAM: 垂直线结构点数 / 全矩阵递归点数
    # 垂直线统计基于全矩阵，所以分母也用全矩阵
    if n_recurrent_total > 0:
        LAM = sum_vert / n_recurrent_total
        LAM = min(LAM, 1.0)
    else:
        LAM = 0.0

    # TT: 捕获时间
    TT = np.mean(vertical_lengths) if len(vertical_lengths) > 0 else 0.0

    return {  # pyright: ignore[reportReturnType]
        'RR': RR,
        'DET': DET,
        'L': L,
        'Lmax': float(Lmax),
        'ENTR': ENTR,
        'LAM': LAM,
        'TT': TT,
    }


# =============================================================================
# RQA 主函数
# =============================================================================


def recurrence_quantification_analysis(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    threshold: float | None = None,
    min_diag_length: int = 2,
    min_vert_length: int = 2,
    axis: int = -1,
    return_dict: bool = False,
) -> Union[np.ndarray, Dict]:
    """递归定量分析 (Recurrence Quantification Analysis, RQA)

    RQA 基于递归图 (Recurrence Plot) 提取定量指标，用于刻画动力系统的复杂度和结构。

    复杂度: O(N² × m)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟，默认自动估计, by default None
    - `threshold`: `float`, 递归阈值，默认 0.1 × std(signal), by default None
    - `min_diag_length`: `int`, 计算 DET 时的最小对角线长度, by default 2
    - `min_vert_length`: `int`, 计算 LAM 时的最小垂直线长度, by default 2
    - `axis`: `int`, 计算轴, by default -1
    - `return_dict`: `bool`, 是否返回字典格式（仅对 1D 输入有效）, by default False

    Returns
    ----------
    - `np.ndarray` 或 `Dict`, RQA 指标数组，顺序为 [RR, DET, L, Lmax, ENTR, LAM, TT]
        - RR: 递归率，系统返回相似状态的频率
        - DET: 确定性，系统行为的可预测程度
        - L: 平均对角线长度，预测时间的度量
        - Lmax: 最大对角线长度，系统稳定性的指示
        - ENTR: 对角线分布熵，系统复杂度
        - LAM: 层叠性，系统停留在特定状态的倾向
        - TT: 捕获时间，平均停留时间
    """
    signal = typing.cast(np.ndarray, ensure_ndarray(signal))

    def _wrapper(
        x: np.ndarray,
        m: int,
        tau: int,
        threshold: float,
        min_diag_length: int,
        min_vert_length: int,
    ) -> np.ndarray:
        _tau = tau if tau is not None else _estimate_tau_autocorr(x)
        _threshold = typing.cast(
            float, threshold if threshold is not None else 0.1 * np.std(x)
        )
        result = _rqa_1d(x, m, _tau, _threshold, min_diag_length, min_vert_length)
        return np.array(
            [
                result['RR'],
                result['DET'],
                result['L'],
                result['Lmax'],
                result['ENTR'],
                result['LAM'],
                result['TT'],
            ]
        )

    result = adapt_1d_func(
        signal, _wrapper, axis, m, tau, threshold, min_diag_length, min_vert_length
    )

    if return_dict and signal.ndim == 1:
        keys = ['RR', 'DET', 'L', 'Lmax', 'ENTR', 'LAM', 'TT']
        return {k: result[i] for i, k in enumerate(keys)}

    return result


# =============================================================================
# 便捷函数
# =============================================================================


def recurrence_rate(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    threshold: float | None = None,
    axis: int = -1,
) -> np.ndarray:
    """计算递归率 (Recurrence Rate, RR)

    递归率表示递归矩阵中递归点的比例，反映系统返回到相似状态的频率。

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟，默认自动估计, by default None
    - `threshold`: `float`, 递归阈值，默认 0.1 * std(signal), by default None
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 递归率
    """
    rqa_result = typing.cast(
        np.ndarray,
        recurrence_quantification_analysis(
            signal, m=m, tau=tau, threshold=threshold, axis=axis
        ),
    )
    if rqa_result.ndim == 1:
        return rqa_result[0]
    return rqa_result[..., 0]


def determinism(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    threshold: float | None = None,
    min_diag_length: int = 2,
    axis: int = -1,
) -> np.ndarray:
    """计算确定性 (Determinism, DET)

    确定性表示属于对角线结构的递归点比例，反映系统行为的可预测程度。
    - 周期信号：DET 接近 1
    - 随机信号：DET 接近 0

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟，默认自动估计, by default None
    - `threshold`: `float`, 递归阈值, by default None
    - `min_diag_length`: `int`, 最小对角线长度, by default 2
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 确定性
    """
    rqa_result = typing.cast(
        np.ndarray,
        recurrence_quantification_analysis(
            signal,
            m=m,
            tau=tau,
            threshold=threshold,
            min_diag_length=min_diag_length,
            axis=axis,
        ),
    )
    if rqa_result.ndim == 1:
        return rqa_result[1]
    return rqa_result[..., 1]


def laminarity(
    signal: np.ndarray,
    m: int = 10,
    tau: int | None = None,
    threshold: float | None = None,
    min_vert_length: int = 2,
    axis: int = -1,
) -> np.ndarray:
    """计算层叠性 (Laminarity, LAM)

    层叠性表示属于垂直线结构的递归点比例，反映系统在某些状态停留的倾向。

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `m`: `int`, 嵌入维度, by default 10
    - `tau`: `int`, 时间延迟, by default None
    - `threshold`: `float`, 递归阈值, by default None
    - `min_vert_length`: `int`, 最小垂直线长度, by default 2
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 层叠性
    """
    rqa_result = typing.cast(
        np.ndarray,
        recurrence_quantification_analysis(
            signal,
            m=m,
            tau=tau,
            threshold=threshold,
            min_vert_length=min_vert_length,
            axis=axis,
        ),
    )
    if rqa_result.ndim == 1:
        return rqa_result[5]
    return rqa_result[..., 5]
