"""Ricerca PRE-REGISTRATA: disaster-cap z-score (z_stop) per la famiglia PAIRS.

Ipotesi pre-registrata: uscita immediata al close della barra se |z| >= z_stop
dopo l'ingresso taglia la coda da structural-break senza toccare i trade normali
(che vivono fra z_exit e z_in).

Griglia PRE-REGISTRATA (unica, completa — NIENTE varianti a posteriori):
  - 5 coppie 1h (config universale n=50 z_in=2.0 z_exit=0.75 max_bars=72):
      z_stop in {3.0, 3.5, 4.0, 5.0}
  - ETH/BTC 15m flat_skip (n=66 z_in=1.674 z_exit=1.0 max_bars=35):
      z_stop in {2.5, 3.0, 3.5, 4.0}
Split: TRAIN = entry prima del 2023-11-01, OOS = dopo (convenzione progetto).

Engine: copia FEDELE di pairs_research.pairs_sim / pairs_sim_flat (stessa
matematica, fee 2 gambe = 2*fee_rt*lev) + parametro z_stop. Causalita': lo z
usato per l'exit alla barra j e' lo stesso z[j] causale (rolling su r[<=j])
gia' usato dall'exit |z|<=z_exit.

REGRESSION-LOCK obbligatorio (eseguito in main, si ferma se fallisce):
  z_stop=None deve riprodurre ESATTAMENTE pairs_sim (ETH/BTC 1h) e
  pairs_sim_flat (ETH/BTC 15m flat_skip): stesso n trade, stesso ret.
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from scripts.analysis.pairs_research import (  # noqa: E402
    FEE_RT, LEV, POS, aligned_ohlc, is_flat_ohlc, pairs_sim, pairs_sim_flat,
)

SPLIT_DT = pd.Timestamp("2023-11-01", tz="UTC")


def pairs_sim_zstop(a, b, tf="1h", n=50, z_in=2.0, z_exit=0.75, max_bars=72,
                    jump_max=0.08, fee_rt=FEE_RT, lev=LEV, pos=POS,
                    z_stop=None, t0=None, t1=None,
                    flat_skip=False, scan_buffer=192):
    """Copia fedele dell'engine pairs (pairs_sim_flat, che con flat_skip=False
    e' identico a pairs_sim — regression-lock in main) + disaster-cap z_stop.

    z_stop: se non None, l'exit si arma anche quando |z[jj]| >= z_stop
    (structural break: lo spread diverge oltre l'ingresso). Stessa convenzione
    causale e stesso fill (close della barra) dell'exit |z|<=z_exit.
    t0/t1: finestra sul timestamp della barra di ENTRY (train/OOS split).
    """
    m = aligned_ohlc(a, b, tf)
    ca, cb = m["close_a"].values, m["close_b"].values
    N = len(ca)
    if flat_skip:
        flat = (is_flat_ohlc(m["open_a"].values, m["high_a"].values, m["low_a"].values, ca)
                | is_flat_ohlc(m["open_b"].values, m["high_b"].values, m["low_b"].values, cb))
    else:
        flat = np.zeros(N, dtype=bool)
    r = np.log(ca / cb)
    dr = np.abs(np.diff(r, prepend=r[0]))
    ma = pd.Series(r).rolling(n).mean().values
    sd = pd.Series(r).rolling(n).std().values
    z = (r - ma) / np.where(sd == 0, np.nan, sd)        # causale: usa r[<=i]
    ts = m["dt"]
    tsv = ts.values                                      # datetime64 per filtro finestra
    t0v = np.datetime64(t0.tz_convert(None)) if t0 is not None else None
    t1v = np.datetime64(t1.tz_convert(None)) if t1 is not None else None
    fee = 2 * fee_rt * lev                               # 2 gambe
    cap = peak = 1000.0; dd = 0.0; last = -1
    trades = wins = n_stop = 0
    rets = []; rets_raw = []
    eq_ts, eq_v = [], []
    kmax = max_bars + (scan_buffer if flat_skip else 0)
    for i in range(n + 1, N - 1):
        if np.isnan(z[i]) or dr[i] > jump_max or i <= last:
            continue
        if t0v is not None and tsv[i] < t0v:
            continue
        if t1v is not None and tsv[i] >= t1v:
            continue
        if z[i] <= -z_in:
            d = 1
        elif z[i] >= z_in:
            d = -1
        else:
            continue
        if flat[i]:
            continue                                   # niente ingresso su barra stale
        # exit: |z|<=z_exit, max_bars, o DISASTER-CAP |z|>=z_stop; con flat_skip
        # l'exit si arma e si esce alla prima barra pulita (live-realizable)
        exit_ready = False; stopped = False; j = i
        for k in range(1, kmax + 1):
            jj = i + k
            if jj >= N:
                j = N - 1; break
            if not exit_ready:
                if z_stop is not None and abs(z[jj]) >= z_stop:
                    exit_ready = True; stopped = True
                elif abs(z[jj]) <= z_exit or k >= max_bars:
                    exit_ready = True
            if exit_ready and not flat[jj]:
                j = jj; break
            j = jj
        retA = (ca[j] - ca[i]) / ca[i]
        retB = (cb[j] - cb[i]) / cb[i]
        ret = (retA - retB) * d * lev - fee             # long A / short B (o viceversa)
        cap = max(cap + cap * pos * ret, 10.0)
        peak = max(peak, cap); dd = max(dd, (peak - cap) / peak)
        trades += 1; wins += ret > 0; n_stop += stopped
        rets.append(ret * pos); rets_raw.append(ret); last = j
        eq_ts.append(ts.iloc[j]); eq_v.append(cap)
    # span temporale della finestra effettiva (per annualizzare lo Sharpe)
    lo = ts.iloc[0] if t0 is None else max(ts.iloc[0], t0)
    hi = ts.iloc[-1] if t1 is None else min(ts.iloc[-1], t1)
    yrs_span = (hi - lo).days / 365.25 or 1
    sharpe = 0.0
    if len(rets) > 1 and np.std(rets) > 0:
        sharpe = float(np.mean(rets) / np.std(rets) * np.sqrt(trades / yrs_span))
    ret_tot = (cap / 1000 - 1) * 100
    worst = min(rets_raw) * 100 if rets_raw else 0.0
    return dict(trades=trades, n_stop=n_stop, win=wins / trades * 100 if trades else 0,
                ret=ret_tot, dd=dd * 100, sharpe=sharpe, worst=worst)


# ----------------------------------------------------------------------------- lock

def regression_lock():
    """z_stop=None deve riprodurre ESATTAVENTE l'engine canonico."""
    ok = True
    # 1h plain vs pairs_sim (config universale live z_exit=0.75)
    ref = pairs_sim("ETH", "BTC", n=50, z_in=2.0, z_exit=0.75, max_bars=72)
    new = pairs_sim_zstop("ETH", "BTC", n=50, z_in=2.0, z_exit=0.75, max_bars=72,
                          z_stop=None, flat_skip=False)
    m1 = (ref["trades"] == new["trades"]) and abs(ref["ret"] - new["ret"]) < 1e-9
    print(f"  LOCK 1h  ETH/BTC vs pairs_sim:      trades {ref['trades']} vs {new['trades']}, "
          f"ret {ref['ret']:+.6f} vs {new['ret']:+.6f} -> {'OK' if m1 else 'FAIL'}")
    ok &= m1
    # 15m flat_skip vs pairs_sim_flat
    ref = pairs_sim_flat("ETH", "BTC", tf="15m", n=66, z_in=1.674, z_exit=1.0,
                         max_bars=35, flat_skip=True)
    new = pairs_sim_zstop("ETH", "BTC", tf="15m", n=66, z_in=1.674, z_exit=1.0,
                          max_bars=35, z_stop=None, flat_skip=True)
    m2 = (ref["trades"] == new["trades"]) and abs(ref["ret"] - new["ret"]) < 1e-9
    print(f"  LOCK 15m ETH/BTC vs pairs_sim_flat: trades {ref['trades']} vs {new['trades']}, "
          f"ret {ref['ret']:+.6f} vs {new['ret']:+.6f} -> {'OK' if m2 else 'FAIL'}")
    ok &= m2
    return ok


# ----------------------------------------------------------------------------- main

PAIRS_1H = [("ETH", "BTC"), ("LTC", "ETH"), ("ADA", "ETH"), ("BTC", "LTC"), ("ETH", "SOL")]
GRID_1H = [None, 3.0, 3.5, 4.0, 5.0]
GRID_15M = [None, 2.5, 3.0, 3.5, 4.0]


def run_cell(a, b, win, z_stop, **kw):
    t0, t1 = (None, SPLIT_DT) if win == "TRAIN" else (SPLIT_DT, None)
    return pairs_sim_zstop(a, b, z_stop=z_stop, t0=t0, t1=t1, **kw)


def main():
    print("=" * 100)
    print("  PAIRS disaster-cap z_stop — ricerca PRE-REGISTRATA (griglia fissa, tutti i risultati)")
    print(f"  split TRAIN < {SPLIT_DT.date()} <= OOS | fee 2 gambe {2*FEE_RT*LEV*100:.2f}% | lev {LEV:.0f}x pos {POS}")
    print("=" * 100)
    print("\nREGRESSION-LOCK (z_stop=None == engine canonico):")
    if not regression_lock():
        print("\n  LOCK FALLITO — STOP."); sys.exit(1)

    hdr = (f"  {'z_stop':>7s} | {'trd':>5s} {'stop':>5s} {'ret%':>9s} {'Shrp':>6s} "
           f"{'DD%':>6s} {'worst%':>8s}")
    for a, b in PAIRS_1H:
        kw = dict(tf="1h", n=50, z_in=2.0, z_exit=0.75, max_bars=72, flat_skip=False)
        print(f"\n{'-'*100}\n  {a}/{b} 1h  (n=50 z_in=2.0 z_exit=0.75 max_bars=72)")
        for win in ("TRAIN", "OOS"):
            print(f"  [{win}]\n{hdr}")
            for zs in GRID_1H:
                r = run_cell(a, b, win, zs, **kw)
                lab = "None" if zs is None else f"{zs:.1f}"
                print(f"  {lab:>7s} | {r['trades']:>5d} {r['n_stop']:>5d} {r['ret']:>+9.1f} "
                      f"{r['sharpe']:>6.2f} {r['dd']:>6.2f} {r['worst']:>+8.2f}")

    a, b = "ETH", "BTC"
    kw = dict(tf="15m", n=66, z_in=1.674, z_exit=1.0, max_bars=35, flat_skip=True)
    print(f"\n{'-'*100}\n  {a}/{b} 15m flat_skip  (n=66 z_in=1.674 z_exit=1.0 max_bars=35)")
    for win in ("TRAIN", "OOS"):
        print(f"  [{win}]\n{hdr}")
        for zs in GRID_15M:
            r = run_cell(a, b, win, zs, **kw)
            lab = "None" if zs is None else f"{zs:.1f}"
            print(f"  {lab:>7s} | {r['trades']:>5d} {r['n_stop']:>5d} {r['ret']:>+9.1f} "
                  f"{r['sharpe']:>6.2f} {r['dd']:>6.2f} {r['worst']:>+8.2f}")


if __name__ == "__main__":
    main()
