"""SKH2_DDKILL — CAUSAL drawdown kill-switch overlay on Skyhook entries.

Family: drawdown kill-switch on entries [DDKILL].

Idea
----
Walk the trade-by-trade REALIZED equity of the V2-winner Skyhook. Track the running peak.
Once standalone DD from the running peak exceeds `dd_kill`, enter a "killed" state and SUPPRESS
new entries until equity recovers within `recover` of the running peak (i.e. DD shrinks back
below `recover`). This is sequential & causal: the kill decision for a new entry at bar i uses
ONLY the equity realized by trades that closed at/before i (busy_until <= i).

Implementation
--------------
1. Build base entries with the winner SkyhookParams.
2. Run backtest_signals -> realized equity path (mark-at-trade-exit, forward-filled).
3. From that equity path compute a per-bar boolean `killed[i]` (causal: peak/DD use eq up to i,
   and eq[i] only changes at a trade-exit bar -> the state at the moment we'd open a new entry
   at i reflects only past closed trades).
4. Null entries where killed -> re-run. Equity changes (suppressed losers/winners during DD),
   so ITERATE to a fixed point (state stabilizes, usually 2-5 iters).
5. Evaluate FULL + HOLD-OUT + fee-sweep + per-year on BOTH assets, marginal-vs-TP01,
   combined-curve max-DD, causality (truncated-prefix recompute of the FINAL entries).

CAUSALITY of the overlay itself
-------------------------------
The base skyhook_entries are already causal (sk.causality). The kill mask at bar i is a function
of equity[0..i], and equity[j] for j<=i only embeds trades whose exit_idx <= i. The mask never
references a future bar. We additionally PROVE it by a truncated-prefix recompute: re-deriving the
final (killed) entries on a data prefix must match the full-run final entries on the overlap.
"""
from __future__ import annotations

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/skyhook")

import numpy as np
import pandas as pd

import skyhooklib as sk
from src.backtest.harness import backtest_signals
from src.strategies.skyhook import SkyhookParams, skyhook_entries

import altlib as al

HOLDOUT = sk.HOLDOUT
FEE = sk.FEE_RT
FEE_SWEEP = (0.0, 0.001, 0.002, 0.003)

# The verified V2 winner from the prior wave.
WINNER = dict(ptn_n=45, sl_atr=2.5, tp_atr=7.0, uscitalong=24, uscitashort=16,
              vola_lo=35.0, vola_hi=95.0, vol_lo=0.0)


def winner_params() -> SkyhookParams:
    return SkyhookParams(**WINNER)


# ---------------------------------------------------------------------------
# Causal DD kill-switch overlay
# ---------------------------------------------------------------------------
def killed_mask_from_equity(equity: np.ndarray, dd_kill: float, recover: float) -> np.ndarray:
    """Per-bar boolean: True = entries SUPPRESSED at this bar.

    State machine over the realized equity path (hysteresis):
      - track running peak.
      - if not killed and DD_from_peak > dd_kill  -> killed.
      - if killed and DD_from_peak <= recover     -> un-killed (recovered).
    Causal: peak[i] and dd[i] use equity[0..i] only.
    """
    n = len(equity)
    killed = np.zeros(n, dtype=bool)
    peak = equity[0]
    state = False
    for i in range(n):
        if equity[i] > peak:
            peak = equity[i]
        dd = (peak - equity[i]) / peak if peak > 0 else 0.0
        if state:
            if dd <= recover:
                state = False
        else:
            if dd > dd_kill:
                state = True
        killed[i] = state
    return killed


def apply_kill(entries: list, killed: np.ndarray) -> list:
    out = list(entries)
    for i in range(len(out)):
        if i < len(killed) and killed[i]:
            out[i] = None
    return out


def ddkill_entries_for_asset(asset: str, p: SkyhookParams, dd_kill: float, recover: float,
                             fee_rt: float = FEE, max_iter: int = 8):
    """Iterate the kill-switch to a fixed point. Returns (final_entries, ltf, n_iters)."""
    ltf, htf = sk.frames(asset)
    base = skyhook_entries(ltf, htf, p)
    cur = base
    prev_killcount = -1
    iters = 0
    for it in range(max_iter):
        m = backtest_signals(ltf, cur, fee_rt=fee_rt, leverage=1.0, asset=asset, tf="230m")
        killed = killed_mask_from_equity(m.equity, dd_kill, recover)
        nxt = apply_kill(base, killed)   # always re-derive from BASE so a recovered state re-enables entries
        iters = it + 1
        kc = int(killed.sum())
        # fixed point: same set of nulled entries as before
        same = all((a is None) == (b is None) for a, b in zip(nxt, cur))
        cur = nxt
        if same and kc == prev_killcount:
            break
        prev_killcount = kc
    return cur, ltf, iters


def _split(eq: np.ndarray, idx: pd.DatetimeIndex, mask: np.ndarray) -> dict:
    e = eq[mask]
    if len(e) < 5:
        return dict(sharpe=0.0, ret=0.0, maxdd=0.0, n=int(len(e)))
    r = np.diff(e) / e[:-1]
    r = r[np.isfinite(r)]
    ix = idx[mask]
    dt = pd.Series(ix).diff().dt.total_seconds().median() or 86400
    bpy = 86400 * 365.25 / dt
    sharpe = float(np.mean(r) / np.std(r) * np.sqrt(bpy)) if len(r) and np.std(r) > 0 else 0.0
    pk = np.maximum.accumulate(e)
    dd = float(np.max((pk - e) / pk))
    return dict(sharpe=round(sharpe, 3), ret=round(float(e[-1] / e[0] - 1), 4),
                maxdd=round(dd, 4), n=int(len(e)))


def study_ddkill(name: str, p: SkyhookParams, dd_kill: float, recover: float):
    per_asset = {}
    fee_ok_all = True
    entries_by_asset = {}
    for a in ("BTC", "ETH"):
        ent, ltf, iters = ddkill_entries_for_asset(a, p, dd_kill, recover)
        entries_by_asset[a] = (ent, ltf)
        m = backtest_signals(ltf, ent, fee_rt=FEE, leverage=1.0, asset=a, tf="230m")
        eq = m.equity
        idx = pd.DatetimeIndex(pd.to_datetime(m.eq_index, utc=True))
        hmask = np.asarray(idx >= HOLDOUT)
        full = dict(sharpe=round(m.sharpe, 3), ret=round(m.net_return, 4), maxdd=round(m.max_dd, 4),
                    n_trades=int(m.n_trades), win_rate=round(m.win_rate, 1))
        hold = _split(eq, idx, hmask)
        sweep = {}
        for f in FEE_SWEEP:
            mf = backtest_signals(ltf, ent, fee_rt=f, leverage=1.0, asset=a, tf="230m")
            sweep[f"{f*100:.2f}%"] = round(mf.sharpe, 3)
        fee_ok_all = fee_ok_all and (sweep["0.30%"] > 0)
        per_asset[a] = dict(full=full, hold=hold,
                            yearly={int(y): round(v, 4) for y, v in m.yearly.items()},
                            fee_sweep=sweep, iters=iters)
    mf = min(per_asset[a]["full"]["sharpe"] for a in per_asset)
    mh = min(per_asset[a]["hold"]["sharpe"] for a in per_asset)
    mt = min(per_asset[a]["full"]["n_trades"] for a in per_asset)
    mdd = max(per_asset[a]["full"]["maxdd"] for a in per_asset)
    grade = "PASS" if (mt >= 20 and mf >= 0.5 and mh >= 0.2 and fee_ok_all) else \
            ("WEAK" if (mt >= 20 and mf >= 0.3 and mh >= 0.0) else "FAIL")
    print(f"\n=== {name} (dd_kill={dd_kill:.0%} recover={recover:.0%}) -> {grade}  "
          f"(minFull={mf:+.2f} minHold={mh:+.2f} minTr={mt} maxDD={mdd*100:.0f}% feeOK={fee_ok_all})")
    for a in per_asset:
        pa = per_asset[a]
        yr = " ".join(f"{y}:{r*100:+.0f}%" for y, r in pa["yearly"].items())
        print(f"  {a}: FULL Sh={pa['full']['sharpe']:+.2f} ret={pa['full']['ret']*100:+.0f}% "
              f"DD={pa['full']['maxdd']*100:.0f}% n={pa['full']['n_trades']} wr={pa['full']['win_rate']:.0f}% "
              f"| HOLD Sh={pa['hold']['sharpe']:+.2f} ret={pa['hold']['ret']*100:+.0f}% DD={pa['hold']['maxdd']*100:.0f}% "
              f"[iters={pa['iters']}]")
        print(f"      fee: " + " ".join(f"{k}={v:+.2f}" for k, v in pa["fee_sweep"].items()))
        print(f"      yr: {yr}")
    return dict(grade=grade, minFull=mf, minHold=mh, minTr=mt, maxDD=mdd, fee_ok=fee_ok_all,
                per_asset=per_asset, entries_by_asset=entries_by_asset,
                dd_kill=dd_kill, recover=recover)


def marginal_ddkill(p: SkyhookParams, dd_kill: float, recover: float):
    def daily(a):
        ent, ltf, _ = ddkill_entries_for_asset(a, p, dd_kill, recover)
        m = backtest_signals(ltf, ent, fee_rt=FEE, leverage=1.0, asset=a, tf="230m")
        s = pd.Series(m.equity, index=pd.DatetimeIndex(pd.to_datetime(m.eq_index, utc=True)))
        return s.resample("1D").last().ffill().pct_change().dropna()
    sb, se = daily("BTC"), daily("ETH")
    J = pd.concat({"BTC": sb, "ETH": se}, axis=1, join="inner").fillna(0.0)
    cand = 0.5 * J["BTC"] + 0.5 * J["ETH"]
    return al.marginal_vs_tp01(cand)


def combined_curve_maxdd(res: dict) -> float:
    """Max-DD of the COMBINED 50/50 BTC+ETH bar-level equity (single standalone curve)."""
    curves = []
    for a in ("BTC", "ETH"):
        ent, ltf = res["entries_by_asset"][a]
        m = backtest_signals(ltf, ent, fee_rt=FEE, leverage=1.0, asset=a, tf="230m")
        s = pd.Series(m.equity, index=pd.DatetimeIndex(pd.to_datetime(m.eq_index, utc=True)))
        curves.append(s.resample("1D").last().ffill().pct_change().fillna(0.0))
    J = pd.concat({"BTC": curves[0], "ETH": curves[1]}, axis=1, join="inner").fillna(0.0)
    r = 0.5 * J["BTC"] + 0.5 * J["ETH"]
    eq = (1.0 + r).cumprod().values
    pk = np.maximum.accumulate(eq)
    return float(np.max((pk - eq) / pk))


# ---------------------------------------------------------------------------
# Causality of the FINAL (killed) entries via truncated-prefix recompute.
# ---------------------------------------------------------------------------
def causality_ddkill(p: SkyhookParams, dd_kill: float, recover: float, asset: str = "BTC",
                     tail: int = 200) -> dict:
    """Re-derive final killed entries on a data PREFIX; they must match the full-run final
    entries on the overlap tail. Proves the kill mask uses no future bar."""
    full_ent, ltf_full = (lambda r: r[:2])(ddkill_entries_for_asset(asset, p, dd_kill, recover))
    n = len(ltf_full)
    ltf, htf = sk.frames(asset)
    bad = 0; checked = 0
    for frac in (0.80, 0.92):
        cut = int(n * frac)
        cut_ts = int(ltf["timestamp"].iloc[cut - 1])
        ltf_sub = ltf.iloc[:cut].reset_index(drop=True)
        htf_sub = htf[htf["timestamp"] <= cut_ts].reset_index(drop=True)
        # re-run the whole kill iteration on the prefix
        base_sub = skyhook_entries(ltf_sub, htf_sub, p)
        cur = base_sub
        for _ in range(8):
            m = backtest_signals(ltf_sub, cur, fee_rt=FEE, leverage=1.0, asset=asset, tf="230m")
            killed = killed_mask_from_equity(m.equity, dd_kill, recover)
            nxt = apply_kill(base_sub, killed)
            same = all((x is None) == (y is None) for x, y in zip(nxt, cur))
            cur = nxt
            if same:
                break
        for i in range(max(0, cut - tail), cut):
            checked += 1
            aE, bE = full_ent[i], cur[i]
            if (aE is None) != (bE is None):
                bad += 1
            elif aE is not None and (aE["dir"] != bE["dir"]
                                     or abs(aE["sl"] - bE["sl"]) > 1e-6
                                     or abs(aE["tp"] - bE["tp"]) > 1e-6):
                bad += 1
    return dict(ok=bool(bad == 0), mismatches=int(bad), checked=int(checked))


def earns_slot_of(res: dict, mg: dict) -> bool:
    return (res["grade"] != "FAIL"
            and mg.get("marginal_verdict") == "ADDS"
            and bool(mg.get("robust_oos"))
            and not bool(mg.get("is_hedge")))


def beats_winner(res: dict, mg: dict, max_dd: float, earns: bool) -> bool:
    w25 = mg.get("blends", {}).get("w25", {}).get("uplift_hold")
    return bool(earns and max_dd < 0.30
                and (w25 is not None and w25 >= 0.55)
                and res["minHold"] >= 0.65)


if __name__ == "__main__":
    p = winner_params()

    print("########## BASELINE (winner, no kill) for reference ##########")
    base_rep = sk.study("WINNER (no kill)", p)
    print(sk.fmt(base_rep))
    base_dd = max(base_rep["per_asset"][a]["full"]["maxdd"] for a in base_rep["per_asset"])
    print(f"  winner standalone maxDD (per-asset max) = {base_dd*100:.0f}%")

    # Grid of (dd_kill, recover). recover < dd_kill (hysteresis: re-enable once DD shrinks back).
    grid = [
        (0.15, 0.10),
        (0.15, 0.12),
        (0.18, 0.12),
        (0.18, 0.15),
        (0.20, 0.15),
        (0.22, 0.16),
        (0.25, 0.18),
        (0.30, 0.22),
    ]

    results = []
    for dd_kill, recover in grid:
        res = study_ddkill(f"DDKILL", p, dd_kill, recover)
        results.append(res)

    print("\n\n########## MARGINAL + combined-DD + earns_slot ##########")
    summary = []
    for res in results:
        mg = marginal_ddkill(p, res["dd_kill"], res["recover"])
        cdd = combined_curve_maxdd(res)
        per_asset_dd = res["maxDD"]
        # standalone max_dd per the brief = max(full.maxdd over BTC & ETH) for the overlay too
        max_dd = per_asset_dd
        earns = earns_slot_of(res, mg)
        bw = beats_winner(res, mg, max_dd, earns)
        w25 = mg.get("blends", {}).get("w25", {}).get("uplift_hold")
        w50 = mg.get("blends", {}).get("w50", {})
        summary.append(dict(dd_kill=res["dd_kill"], recover=res["recover"], grade=res["grade"],
                            minFull=res["minFull"], minHold=res["minHold"], minTr=res["minTr"],
                            per_asset_dd=per_asset_dd, combined_dd=cdd, max_dd=max_dd,
                            corr_full=mg.get("corr_full"), verdict=mg.get("marginal_verdict"),
                            insample=mg.get("has_insample_edge"), hedge=mg.get("is_hedge"),
                            robust=mg.get("robust_oos"), multicut=mg.get("multicut_persistent"),
                            cleanYr=mg.get("clean_year_uplift"), w25=w25, w50=w50,
                            fee_ok=res["fee_ok"], earns=earns, beats=bw, mg=mg, res=res))
        print(f"[dd_kill={res['dd_kill']:.0%} recover={res['recover']:.0%}] grade={res['grade']} "
              f"minFull={res['minFull']:+.2f} minHold={res['minHold']:+.2f} "
              f"perAssetDD={per_asset_dd*100:.0f}% combinedDD={cdd*100:.0f}% "
              f"| corr={mg.get('corr_full')} verdict={mg.get('marginal_verdict')} "
              f"insample={mg.get('has_insample_edge')} hedge={mg.get('is_hedge')} "
              f"robust={mg.get('robust_oos')} multicut={mg.get('multicut_persistent')} "
              f"cleanYr={mg.get('clean_year_uplift')} w25uplift={w25} "
              f"| earns={earns} BEATS={bw}")

    # Pick best honestly: prefer beats_winner; then earns_slot AND a healthy hold-out floor
    # (>=0.65) so we never pick a DD win that kills the hold-out; then lowest per-asset DD;
    # then highest minHold.
    def rank(s):
        healthy = bool(s["earns"]) and (s["minHold"] or -9) >= 0.65
        return (not s["beats"], not healthy, not s["earns"], s["max_dd"], -(s["minHold"] or -9))
    summary.sort(key=rank)
    best = summary[0]
    res = best["res"]; mg = best["mg"]

    # Causality of the FINAL killed entries on the best config, both assets.
    cz_btc = causality_ddkill(p, best["dd_kill"], best["recover"], "BTC")
    cz_eth = causality_ddkill(p, best["dd_kill"], best["recover"], "ETH")
    cz_ok = cz_btc["ok"] and cz_eth["ok"]

    print("\n\n################## BEST CONFIG ##################")
    print(f"config: WINNER + DDKILL(dd_kill={best['dd_kill']:.0%}, recover={best['recover']:.0%})")
    print(f"  minFull        = {best['minFull']:+.3f}")
    print(f"  minHold        = {best['minHold']:+.3f}")
    print(f"  per-asset maxDD= {best['per_asset_dd']*100:.1f}%  (max over BTC&ETH full.maxdd)")
    print(f"  combined  maxDD= {best['combined_dd']*100:.1f}%  (50/50 daily curve)")
    print(f"  n_trades_min   = {best['minTr']}")
    print(f"  fee@0.30%      = {best['fee_ok']}")
    print(f"  causality      = BTC {cz_btc} | ETH {cz_eth} -> ok={cz_ok}")
    print(f"  --- MARGINAL vs TP01 ---")
    print(f"  corr_full           = {mg.get('corr_full')}")
    print(f"  corr_hold           = {mg.get('corr_hold')}")
    print(f"  marginal_verdict    = {mg.get('marginal_verdict')}")
    print(f"  has_insample_edge   = {mg.get('has_insample_edge')}")
    print(f"  is_hedge            = {mg.get('is_hedge')}")
    print(f"  robust_oos          = {mg.get('robust_oos')}")
    print(f"  multicut_persistent = {mg.get('multicut_persistent')}")
    print(f"  clean_year_uplift   = {mg.get('clean_year_uplift')}")
    print(f"  jackknife_min_uplift= {mg.get('jackknife_min_uplift')}")
    print(f"  cand_insample_sharpe= {mg.get('cand_insample_sharpe')}")
    print(f"  blends.w25          = {mg.get('blends', {}).get('w25')}")
    print(f"  blends.w50          = {mg.get('blends', {}).get('w50')}")
    earns = best["earns"]
    print(f"  earns_slot          = {earns}")
    print(f"  BEATS_WINNER        = {best['beats']}")

    # Emit a compact machine-readable line for the orchestrator.
    import json
    w25 = mg.get("blends", {}).get("w25", {}).get("uplift_hold")
    out = dict(
        family="ddkill_entries",
        best_config=dict(base=WINNER, dd_kill=best["dd_kill"], recover=best["recover"]),
        ran_ok=True, grade=res["grade"],
        min_full_sharpe=round(float(best["minFull"]), 3),
        min_hold_sharpe=round(float(best["minHold"]), 3),
        max_dd=round(float(best["max_dd"]), 4),
        combined_dd=round(float(best["combined_dd"]), 4),
        n_trades_min=int(best["minTr"]),
        fee_survives_030=bool(best["fee_ok"]),
        causality_ok=bool(cz_ok),
        marginal_verdict=mg.get("marginal_verdict"),
        has_insample_edge=bool(mg.get("has_insample_edge")),
        is_hedge=bool(mg.get("is_hedge")),
        robust_oos=bool(mg.get("robust_oos")),
        multicut_persistent=bool(mg.get("multicut_persistent")),
        clean_year_uplift=mg.get("clean_year_uplift"),
        corr_full=mg.get("corr_full"),
        blend_w25_uplift_hold=w25,
        earns_slot=bool(earns),
        beats_winner=bool(best["beats"]),
    )
    print("\nRESULT_JSON " + json.dumps(out, default=str))