"""RSK02 — TSMOM long-flat with fast kill-switch on sharp short-horizon drawdown.

IDEA:
  Base signal = TSMOM (multi-horizon momentum: 1m, 3m, 6m) long-flat, vol-targeted (TP01-style).
  Kill-switch: if the position is long AND price has dropped >= `dd_thresh` (e.g. -10%) in the
  last `dd_bars` bars, go flat immediately (hold 0) until momentum re-triggers.

  The kill-switch aims to avoid the worst tail events that TSMOM rides through (sharp crashes).
  It should not improve Sharpe much but should cut max drawdown meaningfully.

Small grid: 2 param sets × 2 TFs = 4 total backtests.
  Config A: dd_thresh=-0.10, dd_bars=5  (10% in 5 bars)
  Config B: dd_thresh=-0.08, dd_bars=3  (8% in 3 bars — tighter)
"""

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np


def tsmom_direction(df) -> np.ndarray:
    """Multi-horizon TSMOM: long if majority of 1m/3m/6m momentum is positive, else flat.
    Causal: uses close[i] returns through i."""
    c = df["close"].values.astype(float)
    bpd = al.bars_per_day(df)

    horizons_days = [21, 63, 126]  # ~1m, 3m, 6m
    signals = []
    for h in horizons_days:
        win = max(2, int(h * bpd))
        # Return over last `win` bars ending at i (causal)
        ret = np.full(len(c), np.nan)
        ret[win:] = c[win:] / c[:-win] - 1.0
        signals.append(np.sign(ret))

    # Vote: positive direction if at least 2 of 3 horizons are positive
    votes = np.nansum(np.stack(signals, axis=0), axis=0)
    direction = np.where(votes > 0, 1.0, 0.0)  # long-flat only
    # Need all 3 to be non-nan (warmup)
    nan_mask = np.any(np.isnan(np.stack(signals, axis=0)), axis=0)
    direction[nan_mask] = 0.0
    return direction


def rolling_drawdown(c: np.ndarray, win: int) -> np.ndarray:
    """Rolling drawdown from the high of the last `win` bars (including current bar i).
    Value at i = (c[i] - max(c[i-win+1:i+1])) / max(...), causal.
    """
    c = c.astype(float)
    n = len(c)
    dd = np.zeros(n)
    # use pandas rolling max (includes current bar)
    import pandas as pd
    rolling_max = pd.Series(c).rolling(win, min_periods=1).max().values
    dd = c / rolling_max - 1.0
    return dd


def make_target(dd_thresh: float, dd_bars: int):
    """Returns a target_fn(df) -> position array."""
    def target_fn(df):
        c = df["close"].values.astype(float)

        # 1. Base TSMOM direction (long or flat)
        direction = tsmom_direction(df)

        # 2. Kill-switch: compute rolling drawdown over dd_bars bars
        rd = rolling_drawdown(c, dd_bars)

        # 3. Kill: if drawdown within last dd_bars is below threshold, go flat
        # We check the minimum drawdown in the last dd_bars window (most severe recent drop)
        import pandas as pd
        # min of rd over last dd_bars: how far price fell from any peak in window
        # Using rolling min of dd to capture worst recent drawdown
        recent_worst_dd = pd.Series(rd).rolling(dd_bars, min_periods=1).min().values
        kill = recent_worst_dd <= dd_thresh  # True = kill signal active

        # Apply kill: override direction to 0 when kill is active
        direction_with_kill = np.where(kill, 0.0, direction)

        # 4. Vol-target the final direction
        tgt = al.vol_target(direction_with_kill, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
        return tgt

    return target_fn


if __name__ == "__main__":
    configs = [
        {"dd_thresh": -0.10, "dd_bars": 5,  "label": "kill10pct-5bar"},
        {"dd_thresh": -0.08, "dd_bars": 3,  "label": "kill08pct-3bar"},
    ]

    best_rep = None
    best_holdout = -999.0

    for cfg in configs:
        name = f"RSK02-{cfg['label']}"
        target_fn = make_target(cfg["dd_thresh"], cfg["dd_bars"])
        rep = al.study_weights(
            name,
            target_fn,
            tfs=("1d", "12h"),
        )
        print(al.fmt(rep))
        print("JSON:", al.as_json(rep))
        print()

        # Track best by holdout sharpe (min across assets)
        ho = rep["verdict"].get("best_holdout_sharpe", -999.0)
        if ho is not None and ho > best_holdout:
            best_holdout = ho
            best_rep = rep

    print("=== BEST CONFIG ===")
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))