"""Mean reversion: z-score of price vs rolling mean with OU half-life estimate and a trend-regime filter."""

from __future__ import annotations

import numpy as np
import pandas as pd

from ktrade.strategy.base import ParamSpec, Strategy, StrategyMeta, log_returns


def ou_half_life(x: pd.Series) -> float:
    """Half-life from an AR(1) fit of Δx on x_{t-1}: HL = -ln2 / ln(1+β). Returns nan when not mean-reverting."""
    x = x.dropna()
    if len(x) < 20:
        return float("nan")
    dx = x.diff().dropna()
    xl = x.shift(1).dropna().loc[dx.index]
    xl_c = xl - xl.mean()
    denom = float((xl_c**2).sum())
    if denom == 0:
        return float("nan")
    beta = float((xl_c * (dx - dx.mean())).sum() / denom)
    if beta >= 0:
        return float("nan")
    return float(-np.log(2) / np.log(1 + beta))


class ZScoreMeanReversion(Strategy):
    meta = StrategyMeta(
        name="zscore_mr",
        label="Z-score mean reversion (OU half-life, regime-filtered)",
        family="mean_reversion",
        hypothesis="Short-horizon deviations of price from its rolling mean revert when the series is locally stationary "
        "(OU half-life short and finite). Enter long when z < -entry (oversold) only if the longer trend is not strongly negative; exit at z > -exit. "
        "Spot-only: the short leg (z > +entry) is not traded. Entries are skipped when the expected reversion is smaller than round-trip costs.",
        experimental=True,
        min_bars=200,
        params=[
            ParamSpec(name="window", type="int", default=48, min=10, max=2000, description="Rolling mean/std window"),
            ParamSpec(name="entry_z", type="float", default=2.0, min=0.5, max=5.0),
            ParamSpec(name="exit_z", type="float", default=0.5, min=0.0, max=3.0),
            ParamSpec(name="max_half_life", type="float", default=48.0, min=1.0, max=5000.0, description="Skip when OU half-life (bars) exceeds this"),
            ParamSpec(name="hl_window", type="int", default=240, min=30, max=5000),
            ParamSpec(name="trend_filter", type="int", default=240, min=0, max=5000, description="EMA span; no entries when price < EMA (0 = off)"),
            ParamSpec(name="cost_bps", type="float", default=60.0, min=0.0, max=500.0, description="Round-trip cost threshold for entry"),
        ],
        failure_modes=["Trending breakdown (catching falling knives)", "Half-life estimate noisy on short windows", "Costs dominate on 1-min bars"],
        references=["Uhlenbeck & Ornstein (1930)", "Chan (2013) Algorithmic Trading, ch. 2"],
    )

    def features(self, df: pd.DataFrame) -> pd.DataFrame:
        p = self.params
        f = pd.DataFrame(index=df.index)
        f["ts"] = df["ts"].values
        f["close"] = df["close"].values
        lc = np.log(df["close"])
        f["ret"] = log_returns(df["close"])
        m = lc.rolling(p["window"], min_periods=p["window"]).mean()
        s = lc.rolling(p["window"], min_periods=p["window"]).std(ddof=1)
        f["z"] = (lc - m) / s.replace(0, np.nan)
        f["half_life"] = lc.rolling(p["hl_window"], min_periods=p["hl_window"]).apply(ou_half_life, raw=False)
        f["expected_rev_bps"] = (-f["z"] * s * 1e4).clip(lower=0)
        if p["trend_filter"] > 0:
            ema = df["close"].ewm(span=p["trend_filter"], adjust=False, min_periods=p["trend_filter"]).mean()
            f["trend_ok"] = (df["close"] >= ema * 0.97).astype(float).where(ema.notna())
        else:
            f["trend_ok"] = 1.0
        return f

    def target(self, feats: pd.DataFrame) -> pd.Series:
        p = self.params
        z = feats["z"].values
        hl = feats["half_life"].values
        ok = feats["trend_ok"].values
        erv = feats["expected_rev_bps"].values
        pos = np.zeros(len(z))
        holding = 0.0
        for i in range(len(z)):
            if np.isnan(z[i]):
                pos[i] = 0.0
                holding = 0.0
                continue
            if holding == 0.0:
                if z[i] < -p["entry_z"] and (np.isnan(hl[i]) is False and hl[i] <= p["max_half_life"]) and ok[i] == 1.0 and erv[i] >= p["cost_bps"]:
                    holding = 1.0
            else:
                if z[i] > -p["exit_z"]:
                    holding = 0.0
            pos[i] = holding
        return pd.Series(pos, index=feats.index)

    def uncertainty(self, feats: pd.DataFrame) -> pd.Series | None:
        hl = feats["half_life"]
        return (hl / self.params["max_half_life"]).clip(0, 1).fillna(1.0)