"""SH01 EXIT POLICY 10 — vol_regime_stop.

Stop CONDIZIONALE al regime di volatilita': lo SL esiste solo quando la vol
sta esplodendo. Razionale: il danno (2020 ETH, crash live 2026-06-05) avviene
in vol-expansion; quando la vol e' normale lo SL taglierebbe winner in
drawdown temporaneo (l'edge SH01 e' nell'asimmetria, win ~50%).

Regime causale: vr[j] = ATR14[j] / SMA100(ATR14)[j]. Nel bar j si guarda
vr[j-1] (dati <= j-1). Se vr[j-1] > r -> SL = entry - d*k1*ATR14[i]
(ATR all'entry = dati <= i). Altrimenti nessuno stop.

  r    in {1.2, 1.5}
  k1   in {1.5, 2.0, 3.0}
  mode in {intrabar, close}

ANTI-LOOK-AHEAD:
  - vr e' un array precomputato module-level (SMA100 causale, no centering).
  - levels(j) usa vr[j-1] e atr14[i] (entry), entrambi <= j-1.
  - mode "close": stop solo se il CLOSE sfonda (stile EXIT-16).
"""
from __future__ import annotations

import sys

sys.path.insert(0, "/opt/docker/PythagorasGoal")

import numpy as np

from scripts.analysis.sh01_exit_lab import (
    ExitPolicy, evaluate, load_sleeves, simulate, OOS_START_MS,
)

_VR_CACHE: dict[int, np.ndarray] = {}


def _vol_ratio(atr14: np.ndarray, win: int = 100) -> np.ndarray:
    """vr[j] = atr14[j] / SMA(atr14, win)[j], causale. NaN dove non definito."""
    key = id(atr14)
    if key in _VR_CACHE:
        return _VR_CACHE[key]
    a = np.asarray(atr14, dtype=float)
    n = len(a)
    sma = np.full(n, np.nan)
    # rolling mean causale (include il bar corrente j: e' OK perche' in levels
    # consumiamo vr[j-1], cioe' dati fino a j-1).
    csum = np.nancumsum(np.where(np.isnan(a), 0.0, a))
    cnt = np.cumsum(~np.isnan(a))
    for j in range(n):
        lo = j - win + 1
        if lo < 0:
            continue
        s = csum[j] - (csum[lo - 1] if lo > 0 else 0.0)
        k = cnt[j] - (cnt[lo - 1] if lo > 0 else 0)
        if k > 0:
            sma[j] = s / k
    vr = np.where((sma > 0) & ~np.isnan(a), a / sma, np.nan)
    _VR_CACHE[key] = vr
    return vr


class VolRegimeStop(ExitPolicy):
    def __init__(self, r: float, k1: float, mode: str):
        self.r = float(r)
        self.k1 = float(k1)
        self.mode = mode
        self.name = f"volreg r{r} k{k1} {mode[:3]}"

    def open_trade(self, ctx: dict, i: int, d: int) -> dict:
        atr_i = ctx["atr14"][i]
        if not np.isfinite(atr_i) or atr_i <= 0:
            atr_i = 0.0
        return {"vr": _vol_ratio(ctx["atr14"]), "atr_i": float(atr_i)}

    def levels(self, ctx: dict, i: int, d: int, j: int, st: dict):
        if j - 1 < 0:
            return None, self.mode
        vr_prev = st["vr"][j - 1]
        if not np.isfinite(vr_prev) or vr_prev <= self.r:
            return None, self.mode          # regime calmo -> nessuno stop
        atr_i = st["atr_i"]
        if atr_i <= 0:
            return None, self.mode
        entry = ctx["close"][i]
        sl = entry - d * self.k1 * atr_i
        return sl, self.mode


# ----------------------------------------------------------------------------- grid

def _fmt(m: dict) -> str:
    return (f"ret={m['ret']:>+7.0f}% dd={m['dd']:>4.0f}% shrp={m['sharpe']:>5.2f} "
            f"worst={m['worst']:>+5.1f}% stop={m['stop_rate']:>4.1f}%")


def main():
    sleeves = load_sleeves()

    # baseline
    print("=== BASELINE (orizzonte puro) ===")
    base = {}
    for a in ("BTC", "ETH"):
        ctx = sleeves[a]
        tr = simulate(ctx, ExitPolicy(), t_hi=OOS_START_MS)
        oo = simulate(ctx, ExitPolicy(), t_lo=OOS_START_MS)
        base[a] = {"train": tr, "oos": oo}
        print(f"  {a} TRAIN {_fmt(tr)}")
        print(f"  {a} OOS   {_fmt(oo)}")

    rs = [1.2, 1.5]
    ks = [1.5, 2.0, 3.0]
    modes = ["intrabar", "close"]

    print("\n=== GRID (TRAIN only) ===")
    grid = {}
    for mode in modes:
        print(f"\n--- mode={mode} ---")
        for r in rs:
            for k1 in ks:
                pol = VolRegimeStop(r, k1, mode)
                btc = simulate(sleeves["BTC"], pol, t_hi=OOS_START_MS)
                eth = simulate(sleeves["ETH"], pol, t_hi=OOS_START_MS)
                grid[(mode, r, k1)] = (btc, eth)
                print(f"  r={r} k1={k1}: BTC {_fmt(btc)} | ETH {_fmt(eth)}")

    # plateau check sul train: cella migliorativa se
    #   ETH sharpe up & dd down & worst less-neg, BTC sharpe>=95% & ret>=80%
    bB_tr = base["BTC"]["train"]; bE_tr = base["ETH"]["train"]
    print("\n=== TRAIN improvement map (cell = ETH sh^ dd_v worst^ AND BTC ok) ===")
    improved = {}
    for key, (btc, eth) in grid.items():
        eth_ok = (eth["sharpe"] > bE_tr["sharpe"] and eth["dd"] < bE_tr["dd"]
                  and eth["worst"] > bE_tr["worst"])
        btc_ok = (btc["sharpe"] >= 0.95 * bB_tr["sharpe"]
                  and btc["ret"] >= 0.80 * bB_tr["ret"])
        improved[key] = eth_ok and btc_ok
        flag = "YES" if improved[key] else " . "
        print(f"  {key}: {flag}  (ethSh {eth['sharpe']:+.2f} vs {bE_tr['sharpe']:+.2f}, "
              f"ethDD {eth['dd']:.0f} vs {bE_tr['dd']:.0f}, ethW {eth['worst']:+.1f} vs {bE_tr['worst']:+.1f}, "
              f"btcSh {btc['sharpe']:.2f} btcRet {btc['ret']:+.0f})")

    n_imp = sum(improved.values())
    print(f"\nTRAIN improving cells: {n_imp}/{len(grid)}")

    # OOS verdict on improving cells (guardato UNA volta)
    print("\n=== OOS verdict (improving train cells) ===")
    for key, ok in improved.items():
        if not ok:
            continue
        mode, r, k1 = key
        pol = VolRegimeStop(r, k1, mode)
        btc = simulate(sleeves["BTC"], pol, t_lo=OOS_START_MS)
        eth = simulate(sleeves["ETH"], pol, t_lo=OOS_START_MS)
        print(f"  {key}: BTC {_fmt(btc)} | ETH {_fmt(eth)}")


if __name__ == "__main__":
    main()
