"""Cointegration analysis for pairs trading.

Provides Engle-Granger and Johansen cointegration tests, rolling analysis,
hedge ratio estimation, and mean-reversion half-life computation.

All functions are pure: data in, results out. No DB access, no API calls.
Requires numpy (core dependency) and statsmodels: ``pip install statsmodels``

Usage::

    from cpz.risk.cointegration import cointegration_test, cointegration_rolling

    result = cointegration_test(prices_a, prices_b, significance=0.05)
    rolling = cointegration_rolling(prices_a, prices_b, window=60)
"""

from __future__ import annotations

import math
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

import numpy as np

try:
    from statsmodels.tsa.stattools import coint
    from statsmodels.tsa.vector_ar.vecm import coint_johansen

    _HAS_SM = True
except ImportError:
    _HAS_SM = False


# ── Models ────────────────────────────────────────────────────────────


class EngleGrangerResult(BaseModel):
    test_statistic: float = 0.0
    p_value: float = 1.0
    critical_values: Dict[str, float] = Field(default_factory=dict)
    is_cointegrated: bool = False


class JohansenResult(BaseModel):
    trace_stat: List[float] = Field(default_factory=list)
    critical_values_90: List[float] = Field(default_factory=list)
    critical_values_95: List[float] = Field(default_factory=list)
    critical_values_99: List[float] = Field(default_factory=list)
    max_eig_stat: List[float] = Field(default_factory=list)
    rank: int = 0


class CointegrationResult(BaseModel):
    """Full cointegration test result."""

    engle_granger: EngleGrangerResult = Field(default_factory=EngleGrangerResult)
    johansen: Optional[JohansenResult] = None
    hedge_ratio: float = 1.0
    half_life: Optional[float] = None
    current_zscore: float = 0.0
    spread_mean: float = 0.0
    spread_std: float = 0.0
    num_observations: int = 0


class RollingCointegrationResult(BaseModel):
    """Rolling cointegration analysis."""

    p_values: List[Optional[float]] = Field(default_factory=list)
    hedge_ratios: List[Optional[float]] = Field(default_factory=list)
    z_scores: List[Optional[float]] = Field(default_factory=list)
    window: int = 60
    num_points: int = 0


# ── Public API ────────────────────────────────────────────────────────


def cointegration_test(
    series_a: List[float],
    series_b: List[float],
    significance: float = 0.05,
) -> CointegrationResult:
    """Run Engle-Granger and Johansen cointegration tests on two price series.

    Args:
        series_a: First price series.
        series_b: Second price series.
        significance: p-value threshold for cointegration (default 0.05).

    Returns:
        CointegrationResult with Engle-Granger, Johansen, hedge ratio,
        half-life, z-score, and spread statistics.

    Raises:
        ImportError: If numpy or statsmodels is not installed.
        ValueError: If fewer than 20 observations are available.
    """
    if not _HAS_SM:
        raise ImportError("Cointegration requires statsmodels: pip install statsmodels")

    a = np.array(series_a, dtype=np.float64)
    b = np.array(series_b, dtype=np.float64)
    n = min(len(a), len(b))
    if n < 20:
        raise ValueError("Need at least 20 observations")
    a, b = a[-n:], b[-n:]

    # Engle-Granger
    eg_stat, eg_pvalue, eg_crit = coint(a, b)

    # Hedge ratio via OLS
    b_mean = np.mean(b)
    a_mean = np.mean(a)
    cov_val = np.mean((b - b_mean) * (a - a_mean))
    var_b = np.mean((b - b_mean) ** 2)
    hedge_ratio = float(cov_val / var_b) if var_b > 1e-15 else 1.0

    # Spread and z-score
    spread = a - hedge_ratio * b
    spread_mean = float(np.mean(spread))
    spread_std = float(np.std(spread))
    current_zscore = float((spread[-1] - spread_mean) / spread_std) if spread_std > 1e-15 else 0.0

    half_life = _compute_half_life(spread)
    johansen_result = _johansen_test(a, b)

    return CointegrationResult(
        engle_granger=EngleGrangerResult(
            test_statistic=round(float(eg_stat), 4),
            p_value=round(float(eg_pvalue), 6),
            critical_values={
                "1%": round(float(eg_crit[0]), 4),
                "5%": round(float(eg_crit[1]), 4),
                "10%": round(float(eg_crit[2]), 4),
            },
            is_cointegrated=float(eg_pvalue) < significance,
        ),
        johansen=johansen_result,
        hedge_ratio=round(hedge_ratio, 6),
        half_life=round(half_life, 2) if half_life and half_life > 0 else None,
        current_zscore=round(current_zscore, 4),
        spread_mean=round(spread_mean, 6),
        spread_std=round(spread_std, 6),
        num_observations=n,
    )


def cointegration_rolling(
    series_a: List[float],
    series_b: List[float],
    window: int = 60,
    significance: float = 0.05,
) -> RollingCointegrationResult:
    """Compute rolling Engle-Granger p-values, hedge ratios, and z-scores.

    Args:
        series_a: First price series.
        series_b: Second price series.
        window: Rolling window size (default 60).
        significance: p-value threshold (default 0.05).

    Returns:
        RollingCointegrationResult with rolling p-values, hedge ratios, z-scores.

    Raises:
        ImportError: If numpy or statsmodels is not installed.
        ValueError: If insufficient data for the given window.
    """
    if not _HAS_SM:
        raise ImportError("Cointegration requires statsmodels: pip install statsmodels")

    a = np.array(series_a, dtype=np.float64)
    b = np.array(series_b, dtype=np.float64)
    n = min(len(a), len(b))
    if n < window + 5:
        raise ValueError(f"Need at least {window + 5} observations for window={window}")
    a, b = a[-n:], b[-n:]

    p_values: List[Optional[float]] = []
    hedge_ratios: List[Optional[float]] = []
    z_scores: List[Optional[float]] = []

    for i in range(n):
        if i < window - 1:
            p_values.append(None)
            hedge_ratios.append(None)
            z_scores.append(None)
            continue

        win_a = a[i - window + 1 : i + 1]
        win_b = b[i - window + 1 : i + 1]

        try:
            _, pval, _ = coint(win_a, win_b)
            p_values.append(round(float(pval), 6))
        except Exception:
            p_values.append(None)

        bm = float(np.mean(win_b))
        am = float(np.mean(win_a))
        cov_val = float(np.mean((win_b - bm) * (win_a - am)))
        var_b = float(np.mean((win_b - bm) ** 2))
        hr = cov_val / var_b if var_b > 1e-15 else 1.0
        hedge_ratios.append(round(hr, 6))

        spread = win_a - hr * win_b
        s_mean = float(np.mean(spread))
        s_std = float(np.std(spread))
        zs = float((spread[-1] - s_mean) / s_std) if s_std > 1e-15 else 0.0
        z_scores.append(round(zs, 4))

    return RollingCointegrationResult(
        p_values=p_values,
        hedge_ratios=hedge_ratios,
        z_scores=z_scores,
        window=window,
        num_points=n,
    )


# ── Internals ─────────────────────────────────────────────────────────


def _compute_half_life(spread: "np.ndarray") -> Optional[float]:
    """Estimate mean-reversion half-life from an AR(1) fit on the spread."""
    spread_lag = spread[:-1]
    spread_ret = np.diff(spread)
    if len(spread_lag) < 3:
        return None
    cov_lr = float(np.mean(
        (spread_lag - np.mean(spread_lag)) * (spread_ret - np.mean(spread_ret))
    ))
    var_lag = float(np.mean((spread_lag - np.mean(spread_lag)) ** 2))
    phi = cov_lr / var_lag if var_lag > 1e-15 else 0.0
    if abs(1 + phi) > 1e-15 and (1 + phi) > 0:
        return -math.log(2) / math.log(abs(1 + phi))
    return None


def _johansen_test(a: "np.ndarray", b: "np.ndarray") -> Optional[JohansenResult]:
    """Run the Johansen trace test on two series."""
    try:
        joh = coint_johansen(np.column_stack([a, b]), det_order=0, k_ar_diff=1)
        return JohansenResult(
            trace_stat=[round(float(x), 4) for x in joh.lr1],
            critical_values_90=[round(float(x), 4) for x in joh.cvt[:, 0]],
            critical_values_95=[round(float(x), 4) for x in joh.cvt[:, 1]],
            critical_values_99=[round(float(x), 4) for x in joh.cvt[:, 2]],
            max_eig_stat=[round(float(x), 4) for x in joh.lr2],
            rank=int(sum(1 for ts, cv in zip(joh.lr1, joh.cvt[:, 1]) if ts > cv)),
        )
    except Exception:
        return None
