"""SKH_R_EXPAND — REGIME variant: VOLATILITY-EXPANSION gate.

Hypothesis (the brief: "enter when vol+volume regime AND breakout coincide"):
Instead of the Chande01 *cycle* band on ATR, define the regime as a genuine VOLATILITY
EXPANSION: trade only when ATR is RISING vs its own moving average (a vol breakout) AND
volume is elevated vs its own moving average. The intuition is that a Donchian breakout that
fires WHILE volatility is expanding on rising participation (volume) is more likely to be a
real move than one that fires inside a quiet/contracting regime (chop, mean-reversion).

Regime definition (HTF, causal):
  vol_expansion = ATR[i] >= k_atr * MA(ATR, w_atr)         (ATR above its own MA -> rising)
  volume_elev   = volume[i] >= k_vol * MA(volume, w_vol)   (participation elevated)
  regime_ok     = vol_expansion AND volume_elev
MA is a CAUSAL rolling mean (uses x[i-w+1..i] inclusive of the current, already-closed bar).
k_atr / k_vol are tunable multipliers (1.0 = "above MA"; >1 = "well above MA").

Pattern/composer/entry/exit reuse the V1 Skyhook building blocks unchanged:
  ptn_n=55 Donchian, sl_atr=2.5, tp_atr=6.0, asymmetric time exits, max 1/day.

Causality: every regime feature uses only x[0..i] (rolling MA, ATR ewm, donchian shift(1)),
INCLUSIVE of the current HTF bar — legit because at HTF close[i] the bar is fully known. The
HTF feature is merged BACKWARD onto LTF on the HTF-close timestamp (<= LTF close). We verify
the truncated-prefix guard ourselves on BOTH assets.
"""
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 import skyhook as S
from src.strategies.skyhook import SkyhookParams, HTF_MIN

HOLDOUT = pd.Timestamp("2025-01-01", tz="UTC")
FEE = sk.FEE_RT


# ---------------------------------------------------------------------------
# Causal rolling MA (inclusive of current, already-closed bar). min_periods enforced.
# ---------------------------------------------------------------------------
def causal_ma(x: np.ndarray, win: int, min_periods: int | None = None) -> np.ndarray:
    mp = win if min_periods is None else min_periods
    return pd.Series(np.asarray(x, float)).rolling(win, min_periods=mp).mean().values


# ---------------------------------------------------------------------------
# HTF feature df: volatility-EXPANSION regime gate + Donchian pattern (V1 pattern reused).
#   regime_ok = (ATR >= k_atr*MA(ATR,w_atr)) AND (volume >= k_vol*MA(volume,w_vol))
# ---------------------------------------------------------------------------
def expand_htf_features(htf: pd.DataFrame, p: SkyhookParams,
                        w_atr: int, k_atr: float,
                        w_vol: int, k_vol: float) -> pd.DataFrame:
    atr_htf = S.atr(htf, p.atr_win)
    vol_htf = htf["volume"].values.astype(float)

    atr_ma = causal_ma(atr_htf, w_atr)
    vol_ma = causal_ma(vol_htf, w_vol)

    # rising-vol = current ATR above k_atr * its own MA ; same for volume.
    # NaN during warmup -> False (no trade until the regime is computable).
    vol_expansion = np.where(np.isfinite(atr_ma) & (atr_ma > 0), atr_htf >= k_atr * atr_ma, False)
    volume_elev = np.where(np.isfinite(vol_ma) & (vol_ma > 0), vol_htf >= k_vol * vol_ma, False)
    regime_ok = vol_expansion & volume_elev

    ptn_long, ptn_short = S.donchian_breakout(htf, p.ptn_n)
    comp_long = regime_ok & ptn_long
    comp_short = regime_ok & ptn_short
    if p.long_only:
        comp_short = np.zeros_like(comp_short, dtype=bool)

    close_ts = htf["timestamp"].astype("int64").values + HTF_MIN * 60 * 1000
    return pd.DataFrame({
        "close_ts": close_ts,
        # store the ratios for diagnostics (not used downstream)
        "buz_vola": np.where(np.isfinite(atr_ma) & (atr_ma > 0), atr_htf / atr_ma, np.nan),
        "buz_volume": np.where(np.isfinite(vol_ma) & (vol_ma > 0), vol_htf / vol_ma, np.nan),
        "comp_long": comp_long.astype(bool), "comp_short": comp_short.astype(bool),
    })


def expand_entries(ltf: pd.DataFrame, htf: pd.DataFrame, p: SkyhookParams,
                   w_atr, k_atr, w_vol, k_vol) -> list:
    """Same entry/exit machinery as S.skyhook_entries, regime from expansion features."""
    feat = expand_htf_features(htf, p, w_atr, k_atr, w_vol, k_vol)
    m = S.merge_htf_to_ltf(ltf, feat)

    c = m["close"].values.astype(float)
    a = S.atr(m, p.ltf_atr_win)
    comp_long = np.nan_to_num(m["comp_long"].values).astype(bool)
    comp_short = np.nan_to_num(m["comp_short"].values).astype(bool)
    days = pd.to_datetime(m["datetime"]).dt.floor("D").values

    entries: list = [None] * len(m)
    count_today: dict = {}
    for i in range(len(m)):
        if not np.isfinite(a[i]) or a[i] <= 0:
            continue
        day = days[i]
        if count_today.get(day, 0) >= p.max_per_day:
            continue
        if comp_long[i]:
            direction, mb = 1, p.uscitalong
        elif comp_short[i]:
            direction, mb = -1, p.uscitashort
        else:
            continue
        if p.exit_mode == "atr":
            sl_off, tp_off = p.sl_atr * a[i], p.tp_atr * a[i]
        else:
            sl_off, tp_off = p.sl_pct * c[i], p.tp_pct * c[i]
        if direction == 1:
            sl, tp = c[i] - sl_off, c[i] + tp_off
        else:
            sl, tp = c[i] + sl_off, c[i] - tp_off
        entries[i] = {"dir": direction, "tp": float(tp), "sl": float(sl), "max_bars": int(mb)}
        count_today[day] = count_today.get(day, 0) + 1
    return entries


# ---------------------------------------------------------------------------
# Eval helpers
# ---------------------------------------------------------------------------
def _split(eq, idx, mask):
    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
    sh = 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(sh, 3), ret=round(float(e[-1] / e[0] - 1), 4), maxdd=round(dd, 4), n=int(len(e)))


def eval_cfg(cfg, p):
    out = {}
    for a in ("BTC", "ETH"):
        ltf, htf = sk.frames(a)
        ent = expand_entries(ltf, htf, p, **cfg)
        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)
        out[a] = dict(full=full, hold=hold, _eq=eq, _idx=idx)
    return out


def summarize(tag, res):
    mf = min(res[a]["full"]["sharpe"] for a in res)
    mh = min(res[a]["hold"]["sharpe"] for a in res)
    mt = min(res[a]["full"]["n_trades"] for a in res)
    mdd = max(res[a]["full"]["maxdd"] for a in res)
    print(f"  [{tag}] minFull={mf:+.2f} minHold={mh:+.2f} minTr={mt} maxDD={mdd*100:.0f}%"
          f" | BTC F{res['BTC']['full']['sharpe']:+.2f}/H{res['BTC']['hold']['sharpe']:+.2f} n{res['BTC']['full']['n_trades']}"
          f" ETH F{res['ETH']['full']['sharpe']:+.2f}/H{res['ETH']['hold']['sharpe']:+.2f} n{res['ETH']['full']['n_trades']}")
    return dict(minFull=mf, minHold=mh, minTr=mt, maxDD=mdd, res=res)


def v1_like_params(**kw):
    base = dict(ptn_n=55, sl_atr=2.5, tp_atr=6.0)  # V1 pattern/exit
    base.update(kw)
    return SkyhookParams(**base)


# ---------------------------------------------------------------------------
# Causality self-check on the structural variant (truncated-prefix guard)
# ---------------------------------------------------------------------------
def check_causality(cfg, p, asset="BTC", tail=150):
    ltf, htf = sk.frames(asset)
    full = expand_entries(ltf, htf, p, **cfg)
    n = len(ltf); bad = 0; checked = 0
    for frac in (0.80, 0.92):
        cut = int(n * frac)
        cut_ts = int(ltf["timestamp"].iloc[cut - 1])
        htf_cut = htf[htf["timestamp"] <= cut_ts].reset_index(drop=True)
        sub = expand_entries(ltf.iloc[:cut].reset_index(drop=True), htf_cut, p, **cfg)
        for i in range(max(0, cut - tail), cut):
            checked += 1
            a, b = full[i], sub[i]
            if (a is None) != (b is None):
                bad += 1
            elif a is not None and (a["dir"] != b["dir"]
                                    or abs(a["sl"] - b["sl"]) > 1e-6
                                    or abs(a["tp"] - b["tp"]) > 1e-6):
                bad += 1
    return dict(ok=bool(bad == 0), mismatches=int(bad), checked=int(checked))


if __name__ == "__main__":
    print("=== SKH_R_EXPAND: volatility-EXPANSION regime (ATR rising vs its MA + volume elevated) ===\n")

    # --- V1 reference (Chande01 regime) ---
    print("--- V1 reference (Chande01 regime, ptn_n=55 sl2.5 tp6 vola35-95 vol0) ---")
    v1 = SkyhookParams(ptn_n=55, sl_atr=2.5, tp_atr=6.0, vola_lo=35, vola_hi=95, vol_lo=0.0)
    v1res = {}
    for a in ("BTC", "ETH"):
        r = sk.run_asset(a, v1, FEE)
        v1res[a] = dict(full=dict(sharpe=r["full"]["sharpe"], n_trades=r["full"]["n_trades"],
                                  maxdd=r["full"]["maxdd"]), hold=dict(sharpe=r["holdout"]["sharpe"]))
    v1_minFull = min(v1res[a]['full']['sharpe'] for a in v1res)
    v1_minHold = min(v1res[a]['hold']['sharpe'] for a in v1res)
    v1_maxDD = max(v1res[a]['full']['maxdd'] for a in v1res)
    print(f"  V1 minFull={v1_minFull:+.2f} minHold={v1_minHold:+.2f} maxDD={v1_maxDD*100:.0f}%"
          f" | BTC F{v1res['BTC']['full']['sharpe']:+.2f}/H{v1res['BTC']['hold']['sharpe']:+.2f}"
          f" ETH F{v1res['ETH']['full']['sharpe']:+.2f}/H{v1res['ETH']['hold']['sharpe']:+.2f}\n")

    p = v1_like_params()  # same pattern/exit as V1; only regime changes

    # --- Sweep: expansion-MA windows + multipliers ---
    # k=1.0 -> "ATR above its own MA" (mild rising). k>1 -> stronger expansion (fewer trades).
    # w_atr/w_vol: lookback for the MA (HTF bars; 690min each). vol elevated mirrored on volume.
    print("--- volatility-EXPANSION sweep ---")
    cfgs = {
        # (w_atr,k_atr) , (w_vol,k_vol)
        "atr20k1.0_vol20k1.0": dict(w_atr=20, k_atr=1.00, w_vol=20, k_vol=1.00),
        "atr20k1.0_vol20k1.2": dict(w_atr=20, k_atr=1.00, w_vol=20, k_vol=1.20),
        "atr20k1.1_vol20k1.0": dict(w_atr=20, k_atr=1.10, w_vol=20, k_vol=1.00),
        "atr20k1.1_vol20k1.2": dict(w_atr=20, k_atr=1.10, w_vol=20, k_vol=1.20),
        "atr10k1.0_vol10k1.0": dict(w_atr=10, k_atr=1.00, w_vol=10, k_vol=1.00),
        "atr10k1.1_vol10k1.2": dict(w_atr=10, k_atr=1.10, w_vol=10, k_vol=1.20),
        "atr30k1.0_vol30k1.0": dict(w_atr=30, k_atr=1.00, w_vol=30, k_vol=1.00),
        "atr20k1.0_volOFF":    dict(w_atr=20, k_atr=1.00, w_vol=20, k_vol=0.00),  # vol gate off (k=0 always true)
        "atr20k1.2_vol20k1.0": dict(w_atr=20, k_atr=1.20, w_vol=20, k_vol=1.00),
    }
    results = {}
    for tag, cfg in cfgs.items():
        results[tag] = (cfg, summarize(tag, eval_cfg(cfg, p)))

    # --- Pick winner by minHold (subject to minFull>=0.5, minTr>=20) ---
    elig = {t: v for t, (c, v) in results.items() if v["minFull"] >= 0.5 and v["minTr"] >= 20}
    print(f"\n--- eligible (minFull>=0.5, minTr>=20): {list(elig.keys())} ---")
    if elig:
        win_tag = max(elig, key=lambda t: elig[t]["minHold"])
    else:
        win_tag = max(results, key=lambda t: results[t][1]["minHold"])
    win_cfg, win_v = results[win_tag][0], results[win_tag][1]
    print(f"\n*** WINNER = {win_tag}  minFull={win_v['minFull']:+.2f} minHold={win_v['minHold']:+.2f}"
          f" minTr={win_v['minTr']} maxDD={win_v['maxDD']*100:.0f}% ***")
    print(f"    cfg = {win_cfg}")

    # --- Causality on winner (BOTH assets) ---
    caus = check_causality(win_cfg, p, "BTC")
    caus_eth = check_causality(win_cfg, p, "ETH")
    print(f"\ncausality(BTC) = {caus}")
    print(f"causality(ETH) = {caus_eth}")

    # --- Fee sweep on winner (both assets, FULL) ---
    print("\n--- fee sweep on winner (FULL sharpe) ---")
    fee_ok_all = True
    for a in ("BTC", "ETH"):
        ltf, htf = sk.frames(a)
        ent_cache = expand_entries(ltf, htf, p, **win_cfg)
        row = []
        for f in (0.0, 0.001, 0.002, 0.003):
            m = backtest_signals(ltf, ent_cache, fee_rt=f, leverage=1.0, asset=a, tf="230m")
            row.append((f, round(m.sharpe, 3)))
        sh030 = dict(row)[0.003]
        fee_ok_all = fee_ok_all and (sh030 > 0)
        print(f"  {a}: " + " ".join(f"{f*100:.2f}%={s:+.2f}" for f, s in row))
    print(f"  fee_survives 0.30%RT (both): {fee_ok_all}")

    # --- Per-year on winner ---
    print("\n--- per-year (winner) ---")
    for a in ("BTC", "ETH"):
        ltf, htf = sk.frames(a)
        ent = expand_entries(ltf, htf, p, **win_cfg)
        m = backtest_signals(ltf, ent, fee_rt=FEE, leverage=1.0, asset=a, tf="230m")
        yr = " ".join(f"{y}:{v*100:+.0f}%" for y, v in m.yearly.items())
        print(f"  {a}: {yr}")

    # --- Marginal vs TP01 on winner ---
    def daily_series(a, p, cfg):
        ltf, htf = sk.frames(a)
        ent = expand_entries(ltf, htf, p, **cfg)
        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()

    import altlib as al
    sb = daily_series("BTC", p, win_cfg); se = daily_series("ETH", p, win_cfg)
    J = pd.concat({"BTC": sb, "ETH": se}, axis=1, join="inner").fillna(0.0)
    cand_daily = 0.5 * J["BTC"] + 0.5 * J["ETH"]
    marg = al.marginal_vs_tp01(cand_daily)
    print("\n--- marginal vs TP01 (winner) ---")
    print(f"  corr_full={marg.get('corr_full')} corr_hold={marg.get('corr_hold')}")
    print(f"  blend w25 uplift_hold={marg.get('blends',{}).get('w25',{}).get('uplift_hold')}")
    print(f"  marginal_verdict={marg.get('marginal_verdict')} robust_oos={marg.get('robust_oos')}")
    print(f"  has_insample_edge={marg.get('has_insample_edge')} is_hedge={marg.get('is_hedge')}")
    print(f"  clean_year_uplift={marg.get('clean_year_uplift')} jackknife_min_uplift={marg.get('jackknife_min_uplift')}")
    print(f"  multicut_persistent={marg.get('multicut_persistent')}")

    # --- Final verdict echo ---
    print("\n=== SUMMARY ===")
    print(f"V1:            minFull={v1_minFull:+.2f} minHold={v1_minHold:+.2f} maxDD={v1_maxDD*100:.0f}%")
    print(f"EXPAND {win_tag}: minFull={win_v['minFull']:+.2f} minHold={win_v['minHold']:+.2f}"
          f" maxDD={win_v['maxDD']*100:.0f}% causalityOK={caus['ok'] and caus_eth['ok']} feeOK={fee_ok_all}")
    beats = (win_v['minHold'] > v1_minHold) and win_v['minFull'] >= 0.5 and fee_ok_all
    print(f"BEATS V1 HOLD-OUT: {beats}")
