"""
Cointegration-based pair discovery with Ornstein-Uhlenbeck
half-life filtering.

Provides tools to:
  1. Test pairs for cointegration (Engle-Granger two-step)
  2. Estimate OU half-life of the residual spread
  3. Filter candidate pairs by minimum half-life
  4. Rank pairs by mean-reversion strength (high ADF stat, low half-life)

All implemented in pure NumPy — no statsmodels dependency.

Design decisions:
  - Critical values for ADF test are hardcoded (MacKinnon 1994 tables)
    → avoids importing statsmodels.
  - Both 1% and 5% significance levels supported.
  - Half-life computed via OLS on the AR(1) of the residual.
  - Minimum observations: 100 for cointegration test (avoid spurious results).
  - Sector constraint: optional list of ticker prefixes (e.g., "ETH", "BTC").
"""

from __future__ import annotations

import numpy as np
from typing import Optional


# ── MacKinnon (1994) critical values for ADF test ────────────
# Table for Case 2: regression with intercept, no trend
# Rows: sample sizes (25, 50, 100, 250, 500, ∞)
# Cols: significance levels (1%, 5%, 10%)

_MACKINNON_CASE2 = np.array([
    [-3.75, -3.00, -2.63],  # N=25
    [-3.58, -2.93, -2.60],  # N=50
    [-3.51, -2.89, -2.58],  # N=100
    [-3.46, -2.88, -2.57],  # N=250
    [-3.44, -2.87, -2.57],  # N=500
    [-3.43, -2.86, -2.57],  # N=∞
])

_MACKINNON_N_SIZES = np.array([25, 50, 100, 250, 500, 999999])


def adf_critical_value(n_obs: int, sig: float = 0.05) -> float:
    """Return ADF critical value for given sample size and significance."""
    col = 0 if sig <= 0.01 else 1 if sig <= 0.05 else 2
    idx = np.searchsorted(_MACKINNON_N_SIZES, n_obs, side="right") - 1
    idx = max(0, min(idx, len(_MACKINNON_N_SIZES) - 1))
    return float(_MACKINNON_CASE2[idx, col])


def adf_test(residuals: np.ndarray, sig: float = 0.05) -> dict:
    """
    Augmented Dickey-Fuller test (no lags).

    Tests H₀: unit root (not mean-reverting) vs H₁: stationary.

    Args:
        residuals: 1-D array of OLS residuals from cointegrating regression.
        sig: Significance level (0.01 or 0.05).

    Returns:
        dict with keys: statistic, critical_value, is_stationary, p_value_approx.
    """
    n = len(residuals)
    if n < 20:
        return {"statistic": 0.0, "critical_value": 0.0, "is_stationary": False, "p_value_approx": 1.0}

    dy = np.diff(residuals)
    y_lag = residuals[:-1]

    # OLS: Δyₜ = γ yₜ₋₁ + εₜ
    X = y_lag.reshape(-1, 1)
    Y = dy.reshape(-1, 1)

    # γ = (XᵀX)⁻¹ XᵀY
    XtX = X.T @ X
    if XtX[0, 0] < 1e-12:
        return {"statistic": 0.0, "critical_value": 0.0, "is_stationary": False, "p_value_approx": 1.0}

    gamma = float((np.linalg.inv(XtX) @ X.T @ Y)[0, 0])
    residuals_ols = Y.flatten() - gamma * X.flatten()
    se = np.std(residuals_ols, ddof=1)
    t_stat = gamma / se if se > 1e-12 else 0.0

    crit = adf_critical_value(n, sig)
    is_stat = t_stat < crit

    # Rough p-value approximation
    p_val = max(0.0, min(1.0, 1.0 / (1.0 + np.exp(-(abs(t_stat) - 2.0)))))

    return {
        "statistic": round(t_stat, 4),
        "critical_value": round(crit, 4),
        "is_stationary": is_stat,
        "p_value_approx": round(p_val, 4),
    }


def estimate_half_life(spread: np.ndarray) -> float:
    """
    Estimate the Ornstein-Uhlenbeck half-life of a spread series.

    Model: dsₜ = θ (μ - sₜ) dt + σ dWₜ

    Half-life = ln(2) / θ

    Implementation:
        Discretize and run OLS on: sₜ₊₁ - sₜ = a + b sₜ + εₜ
        Then θ = -b, half-life = ln(2) / θ.

    Args:
        spread: 1-D array of spread values.

    Returns:
        Half-life in number of periods. Returns inf if not mean-reverting.
    """
    n = len(spread)
    if n < 20:
        return float("inf")

    s = spread
    ds = np.diff(s)
    s_lag = s[:-1]

    # OLS: ds[t] = a + b * s[t-1]
    X = np.column_stack([np.ones(len(s_lag)), s_lag])
    Y = ds

    try:
        coeff = np.linalg.lstsq(X, Y, rcond=None)[0]
    except np.linalg.LinAlgError:
        return float("inf")

    b = coeff[1]  # mean-reversion speed (negative → mean-reverting)

    if b >= 0:
        return float("inf")  # Not mean-reverting

    theta = -b
    half_life = np.log(2) / theta if theta > 1e-10 else float("inf")

    return float(half_life)


def test_pair(X: np.ndarray, Y: np.ndarray, sig: float = 0.05) -> dict:
    """
    Full cointegration + half-life test for a candidate pair.

    Engle-Granger two-step:
        1. Regress Y on X: Y = α + β X + ε
        2. Test ε for stationarity (ADF)
        3. Estimate half-life of ε

    Args:
        X: Price series of asset X (independent).
        Y: Price series of asset Y (dependent).
        sig: ADF significance level.

    Returns:
        dict with:
            alpha, beta (hedge ratio), adf_stat, adf_crit,
            is_cointegrated, half_life, half_life_days,
            spread, spread_std, correlation
    """
    n = min(len(X), len(Y))
    if n < 100:
        return {"is_cointegrated": False, "half_life": float("inf"), "reason": "insufficient_data"}

    x = np.array(X[-n:])
    y = np.array(Y[-n:])

    # Step 1: OLS regression
    X_mat = np.column_stack([np.ones(n), x])
    try:
        coeff = np.linalg.lstsq(X_mat, y, rcond=None)[0]
    except np.linalg.LinAlgError:
        return {"is_cointegrated": False, "half_life": float("inf"), "reason": "lstsq_failed"}

    alpha, beta = float(coeff[0]), float(coeff[1])

    # Step 2: Residuals
    residuals = y - (alpha + beta * x)

    # ADF test on residuals
    adf = adf_test(residuals, sig=sig)

    # Step 3: Half-life
    hl = estimate_half_life(residuals)

    return {
        "alpha": round(alpha, 6),
        "beta": round(beta, 6),
        "adf_stat": adf["statistic"],
        "adf_crit": adf["critical_value"],
        "is_cointegrated": adf["is_stationary"],
        "half_life": round(hl, 2),
        "spread": residuals,
        "spread_std": round(float(np.std(residuals)), 6),
        "correlation": round(float(np.corrcoef(x, y)[0, 1]), 4),
    }


def discover_pairs(
    price_data: dict[str, np.ndarray],
    sector_constraint: Optional[str] = None,
    min_half_life: float = 1.0,
    max_half_life: float = 20.0,
    sig_level: float = 0.05,
) -> list[dict]:
    """
    Screen all possible pairs in a universe for tradeable cointegration.

    Filters:
        1. ADF test passes at given significance level
        2. Half-life between min_half_life and max_half_life (periods)
        3. Optional sector constraint (ticker prefix match)

    Args:
        price_data: {ticker: price_array} mapping.
        sector_constraint: If set, only pairs where both tickers share this prefix.
        min_half_life: Minimum half-life in periods.
        max_half_life: Maximum half-life in periods.
        sig_level: ADF significance level.

    Returns:
        List of dicts, sorted by half-life (ascending — faster mean reversion first).
        Each dict has: pair, alpha, beta, half_life, adf_stat, spread_std, correlation.
    """
    tickers = sorted(price_data.keys())
    results: list[dict] = []

    for i in range(len(tickers)):
        for j in range(i + 1, len(tickers)):
            t1, t2 = tickers[i], tickers[j]

            # Sector constraint
            if sector_constraint:
                if not (t1.startswith(sector_constraint) and t2.startswith(sector_constraint)):
                    continue

            X = price_data[t1]
            Y = price_data[t2]

            test = test_pair(X, Y, sig=sig_level)
            if test["is_cointegrated"] and min_half_life <= test["half_life"] <= max_half_life:
                results.append({
                    "pair": (t1, t2),
                    "X_ticker": t1,
                    "Y_ticker": t2,
                    "alpha": test["alpha"],
                    "beta": test["beta"],
                    "half_life": test["half_life"],
                    "adf_stat": test["adf_stat"],
                    "spread_std": test["spread_std"],
                    "correlation": test["correlation"],
                })

    # Sort by half-life (faster mean reversion = better)
    results.sort(key=lambda r: r["half_life"])
    return results


def compute_rolling_ols_hedge(
    X: np.ndarray,
    Y: np.ndarray,
    window: int = 60,
) -> np.ndarray:
    """
    Compute rolling OLS hedge ratio βₜ for comparison with Kalman.

    Uses expanding window OLS up to the specified lookback.

    Args:
        X, Y: Price series.
        window: Lookback window in periods.

    Returns:
        1-D array of β values (same length as inputs).
    """
    n = len(X)
    betas = np.full(n, np.nan)
    for t in range(window, n):
        x_win = X[t - window:t]
        y_win = Y[t - window:t]
        X_mat = np.column_stack([np.ones(len(x_win)), x_win])
        try:
            coeff = np.linalg.lstsq(X_mat, y_win, rcond=None)[0]
            betas[t] = coeff[1]
        except np.linalg.LinAlgError:
            betas[t] = np.nan
    return betas