"""Hurst 指数

Hurst 指数用于评估时间序列的长期记忆性和分形特性。
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func


# =============================================================================
# Hurst 指数 (Hurst Exponent)
# =============================================================================


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:
                    rs_block_ratios[j] = 0.0  # 0/0 case, treat as 0
                else:
                    rs_block_ratios[j] = np.nan  # R > 0, S = 0
            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)

    Hurst 指数的解释：
    - H = 0.5: 随机游走 (无记忆性)
    - H > 0.5: 持续性序列 (正相关)
    - H < 0.5: 反持续性序列 (负相关)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `max_lag`: `int`, 最大滞后时间 (inclusive), by default 20
    - `axis`: `int`, 计算轴, by default -1

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

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