"""Minimal stochastic-process feature computation — OU mean-reversion + Hurst.

These are the features that beat hand-rolled TA in this repo's 50-ETF runs.
Compute them per symbol on a rolling window ENDING at each day (never let them
see test data at fit time — see the HMM/GARCH note in SKILL.md).

Reference repo impl: tac-qlib/examples/sp_features.py (full 55-feature set:
OU, HMM, jump, HARRV, trend, GARCH, Hurst, path signatures, entropy, catch22).
"""

from __future__ import annotations

import numpy as np
import pandas as pd

#: rolling window for feature computation (days)
LOOKBACK = 250


def compute_ou_features(close: pd.Series) -> pd.DataFrame:
    """Ornstein-Uhlenbeck fit: theta (reversion speed), sigma (vol), residual z.

    OU: dx_t = theta (mu - x_t) dt + sigma dW_t  (theta is the mean-reversion
    speed; higher = faster reversion = tradable mean-reversion signal).

    Rolling OLS of dx on lagged log-price gives theta = -b (reversion speed);
    sigma is the residual std. Vectorized via rolling cov/var.
    """
    logp = np.log(close)
    dx = logp.diff()
    x_prev = logp.shift(1)
    df = pd.DataFrame({"dx": dx, "x": x_prev})

    out = pd.DataFrame(index=close.index, dtype=float)
    cov = df["dx"].rolling(LOOKBACK, min_periods=30).cov(df["x"])
    var = df["x"].rolling(LOOKBACK, min_periods=30).var()
    theta = (-cov / var).rename("sp_ou_theta")
    out["sp_ou_theta"] = theta
    out["sp_ou_sigma"] = df["dx"].rolling(LOOKBACK, min_periods=30).std()
    # standardized residual z = (x - mu) / sigma of the fitted process
    mu = df["x"].rolling(LOOKBACK, min_periods=30).mean()
    scale = np.sqrt(np.clip(1 / (2 * theta + 1e-9), 0, None))
    out["sp_ou_zscore"] = (df["x"] - mu) / (out["sp_ou_sigma"] * scale)
    return out


def compute_hurst(close: pd.Series, lookback: int = 100) -> pd.Series:
    """Rolling Hurst exponent via rescaled range (R/S). H>0.5 = trending."""
    def _hurst(x: np.ndarray) -> float:
        if len(x) < 20:
            return np.nan
        lags = range(2, min(len(x) // 2, 50))
        tau = []
        for lag in lags:
            diff = x[lag:] - x[:-lag]
            tau.append(np.sqrt(np.std(diff)))
        tau = np.array(tau)
        lags = np.array(lags, dtype=float)
        poly = np.polyfit(np.log(lags), np.log(tau), 1)
        return float(poly[0])

    return close.rolling(lookback, min_periods=20).apply(lambda w: _hurst(w.to_numpy()), raw=False).rename(
        "sp_hurst_exponent"
    )


def build_sp_features(bars: pd.DataFrame) -> pd.DataFrame:
    """bars: lake 1d bars indexed by (datetime, instrument) or a symbol frame."""
    if isinstance(bars.index, pd.MultiIndex):
        frames = []
        for inst, sub in bars.groupby(level=1):
            close = sub.droplevel(1)["close"]
            feats = pd.concat([compute_ou_features(close), compute_hurst(close)], axis=1)
            feats["instrument"] = inst
            frames.append(feats.reset_index())
        out = pd.concat(frames).set_index(["datetime", "instrument"])
    else:
        close = bars["close"]
        out = pd.concat([compute_ou_features(close), compute_hurst(close)], axis=1)
    return out