import numpy as np
import statsmodels.api as sm
from statsmodels.tsa.stattools import adfuller, coint

def calculate_hurst(ts: np.ndarray) -> float:
    """使用 R/S 分析计算 Hurst 指数"""
    try:
        lags = range(2, min(100, len(ts) // 2))
        tau = [np.std(np.subtract(ts[lag:], ts[:-lag])) for lag in lags]
        # 避免 log(0)
        tau = [t for t in tau if t > 0]
        lags = list(lags)[:len(tau)]
        if len(lags) < 2: return 0.5
        poly = np.polyfit(np.log(lags), np.log(tau), 1)
        return poly[0] * 2.0
    except Exception:
        return 0.5 # 计算失败返回随机游走假设

def calculate_annualized_vol(returns: np.ndarray, periods_per_year: int) -> float:
    """计算年化波动率"""
    if len(returns) < 2: return 0.0
    return np.std(returns) * np.sqrt(periods_per_year)

def check_adf_stationarity(ts: np.ndarray, max_pvalue: float) -> bool:
    """ADF 平稳性检验"""
    try:
        # 剔除 NaN
        ts = ts[~np.isnan(ts)]
        if len(ts) < 20: return False
        result = adfuller(ts, autolag='AIC')
        return result[1] <= max_pvalue
    except Exception:
        return False

def check_cointegration(ts1: np.ndarray, ts2: np.ndarray, max_pvalue: float) -> bool:
    """Engle-Granger 协整检验"""
    try:
        # 对齐长度并剔除 NaN
        min_len = min(len(ts1), len(ts2))
        ts1, ts2 = ts1[-min_len:], ts2[-min_len:]
        mask = ~np.isnan(ts1) & ~np.isnan(ts2)
        ts1, ts2 = ts1[mask], ts2[mask]
        
        if len(ts1) < 50: return False
        score, pvalue, _ = coint(ts1, ts2)
        return pvalue <= max_pvalue
    except Exception:
        return False