"""_idmom02.py - IDMOM02: market intraday momentum on US100 / NQ, last half hour.
Prereg: IDMOM02_PREREG.md (frozen and pushed before this run). Tested by: Claude (2026-10-10).
signal = sign(close 15:29 bar / previous regular-session close - 1); entry open 15:30; exit close 15:59;
stop and risk unit = 1.5 x ATR14 of RTH M15 bars at the 15:15 bar (features() imported unchanged).
Run:  python -X utf8 prop_eur2k_lab/pillar2_hunt/_idmom02.py
"""
import math
import os
import sys

import numpy as np
import pandas as pd
import polars as pl

LAB = r"<repo>\prop_eur2k_lab"
sys.path.insert(0, LAB)
import _bar_boundary_audit_20261009 as B  # noqa: E402
# [withheld: refers to the lab's private live strategy]

NY = "America/New_York"
OUT = os.path.join(LAB, "pillar2_hunt")
LSE = r"<repo>\quantlab\data\inbox\lse\NQ_F_1m_2016-06-01_2026-10-09.parquet"
DBN = r"D:\quant_data\mega_futures\extracted"
CFD_PTS, FUT_PTS = 1.70, 0.50 + 2 * 2.25 / 20.0          # round trip, at the reference price
RNG = np.random.default_rng(20261010)


def tstat(x):
    x = np.asarray(x, float)
    x = x[~np.isnan(x)]
    return float(x.mean() / x.std(ddof=1) * math.sqrt(len(x))) if len(x) > 2 and x.std(ddof=1) > 0 else float("nan")


def build(m, prev_close=None):
    """m: tz-aware America/New_York M1 bars stamped at minute START, columns open/high/low/close. One row per usable day."""
    m = m[m.index.dayofweek < 5]
    mod = m.index.hour * 60 + m.index.minute
    rth = m[(mod >= 570) & (mod <= 959)]
    rmod = np.asarray(rth.index.hour * 60 + rth.index.minute)
    day = rth.index.normalize().tz_localize(None)
    assert not rth.index.duplicated().any()

    def at(col, mm):
        k = rmod == mm
        return pd.Series(rth[col].to_numpy()[k], index=day[k])

    D = pd.DataFrame({"n": pd.Series(1, index=day).groupby(level=0).sum()})
    last = pd.Series(rth["close"].to_numpy(), index=day).groupby(level=0).last()
    D["prev_close"] = last.shift(1) if prev_close is None else prev_close.reindex(D.index)
    D["gap_days"] = pd.Series(D.index, index=D.index).diff().dt.days
    for name, col, mm in [("c1529", "close", 929), ("o1530", "open", 930), ("c1559", "close", 959),
                          ("c1544", "close", 944), ("c0959", "close", 599)]:
        D[name] = at(col, mm)
    for tag, hi_mm in [("", 959), ("_45", 944)]:
        k = (rmod >= 930) & (rmod <= hi_mm)
        D["lo" + tag] = pd.Series(rth["low"].to_numpy()[k], index=day[k]).groupby(level=0).min()
        D["hi" + tag] = pd.Series(rth["high"].to_numpy()[k], index=day[k]).groupby(level=0).max()
    # [withheld: refers to the lab's private live strategy]
    a = features(g)["atr"]
    a = a[(a.index.hour == 15) & (a.index.minute == 15)]
    D["risk"] = 1.5 * pd.Series(a.to_numpy(), index=a.index.normalize().tz_localize(None))
    D = D[(D["n"] >= 360) & (D["gap_days"] <= 5)].dropna(subset=["prev_close", "c1529", "o1530", "c1559", "risk"])
    D = D[D["risk"] > 0]
    D["rod"] = D["c1529"] / D["prev_close"] - 1
    D["sig"] = np.sign(D["rod"])
    D = D[D["sig"] != 0].copy()
    pts = D["c1559"] - D["o1530"]
    D["R_long"] = np.where(D["lo"] <= D["o1530"] - D["risk"], -1.0, pts / D["risk"])
    D["R_short"] = np.where(D["hi"] >= D["o1530"] + D["risk"], -1.0, -pts / D["risk"])
    D["R"] = np.where(D["sig"] > 0, D["R_long"], D["R_short"])
    D["stopped"] = D["R"] == -1.0
    D["bp"] = D["sig"] * (D["c1559"] / D["o1530"] - 1) * 1e4               # no stop, for intuition
    p45 = D["c1544"] - D["o1530"]
    D["R45"] = np.where(D["sig"] > 0, np.where(D["lo_45"] <= D["o1530"] - D["risk"], -1.0, p45 / D["risk"]),
                        np.where(D["hi_45"] >= D["o1530"] + D["risk"], -1.0, -p45 / D["risk"]))
    D["sig_gao"] = np.sign(D["c0959"] / D["prev_close"] - 1)
    D["year"] = D.index.year
    return D


def add_cost(D, pts, ref):
    D = D.copy()
    D["cost"] = pts * D["o1530"] / ref / D["risk"]
    D["net"], D["net2"] = D["R"] - D["cost"], D["R"] - 2 * D["cost"]
    return D


def row(lab, D, col="net"):
    r = D[col].dropna().to_numpy(float)
    if len(r) < 3:
        return f"  {lab:<34s} n={len(r):5d}"
    return (f"  {lab:<34s} n={len(r):5d}  mean={r.mean():+.4f}R  t={tstat(r):+5.2f}  win={100 * (r > 0).mean():4.1f}%  "
            f"sum={r.sum():+7.1f}R")


def dbn_m1(sym):
    """Databento 1s -> RTH M1 of the day's front contract (by volume); previous close from the SAME contract."""
    lf = (pl.scan_parquet(os.path.join(DBN, f"{sym}_ohlcv-1s_2010-2018.parquet"))
          .filter(~pl.col("symbol").str.contains("-"))
          .filter(pl.col("ts_event") >= pl.datetime(2012, 12, 1, time_zone="UTC"))
          .with_columns(pl.col("ts_event").dt.convert_time_zone(NY).alias("ts"))
          .filter((pl.col("ts").dt.time() >= pl.time(9, 30)) & (pl.col("ts").dt.time() < pl.time(16, 0)))
          .with_columns(pl.col("ts").dt.date().alias("d")))
    df = lf.select(["ts", "d", "symbol", "open", "high", "low", "close", "volume"]).collect().sort("ts")
    per = df.group_by(["d", "symbol"], maintain_order=True).agg(pl.col("volume").sum().alias("v"), pl.col("close").last().alias("c"))
    front = per.sort(["d", "v"], descending=[False, True]).group_by("d", maintain_order=True).first().sort("d")
    days = front["d"].to_list()
    last = {(r[0], r[1]): r[3] for r in per.iter_rows()}
    prev = {d: last.get((days[i - 1], s)) for i, (d, s) in enumerate(zip(days, front["symbol"].to_list())) if i > 0}
    df = df.join(front.select(["d", "symbol"]), on=["d", "symbol"], how="inner").sort("ts")
    m1 = df.group_by_dynamic("ts", every="1m").agg(pl.col("open").first(), pl.col("high").max(), pl.col("low").min(), pl.col("close").last())
    p = m1.to_pandas().set_index("ts").sort_index()
    p.index = p.index.tz_convert(NY)
    pc = pd.Series(prev, dtype=float)
    pc.index = pd.to_datetime(pc.index)
    return p, pc


def main():
    cfd = B.load_cfd()
    ref = float(cfd["close"].iloc[-5000:].median())
    print(f"reference price {ref:,.0f} | CFD cost {CFD_PTS} pt, futures cost {FUT_PTS:.3f} pt round trip, scaled by price\n")

    print("1. PRIMARY: NAS100 CFD, true clock (file stamps - 1 min), CFD cost")
    books = {ph: add_cost(build(B.shifted(cfd, ph)), CFD_PTS, ref) for ph in range(-3, 4)}
    D = books[-1]
    D.to_csv(os.path.join(OUT, "_idmom02_trades_cfd.csv"))
    print(f"  days {D.index.min().date()} -> {D.index.max().date()} | stop hit {100 * D['stopped'].mean():.1f}% | "
          f"risk median {1e4 * (D['risk'] / D['o1530']).median():.0f} bp of price | cost median {D['cost'].median():.4f}R | SD of R {D['R'].std():.3f}")
    print(row("gross", D, "R"))
    print(row("net CFD 1x", D))
    print(row("net CFD 2x", D, "net2"))
    print(f"  mean move with the signal, no stop: {D['bp'].mean():+.2f} bp (t {tstat(D['bp']):+.2f}); cost {1e4 * CFD_PTS / ref:.2f} bp")
    print(row("2005-2015 net", D[D.year <= 2015]))
    print(row("2016-2026 net", D[D.year >= 2016]))
    print(row("2019-2026 net (post-publication)", D[D.year >= 2019]))
    print(row("signal long days, net", D[D.sig > 0]))
    print(row("signal short days, net", D[D.sig < 0]))
    print(row("report-only: exit 15:45, net", D.assign(x=D["R45"] - D["cost"]), "x"))
    gk = D["sig_gao"] != 0
    gao = np.where(D["sig_gao"] > 0, D["R_long"], D["R_short"])[gk] - D["cost"][gk]
    print(row("report-only: Gao 10:00 signal, net", pd.DataFrame({"net": gao})))
    print("  by year (net 1x, sum R / t): " + "  ".join(f"{y}:{g['net'].sum():+.0f}/{tstat(g['net']):+.1f}" for y, g in D.groupby("year")))
    print("  phases (net 1x t): " + "  ".join(f"{ph:+d}m {tstat(books[ph]['net']):+.2f}" for ph in range(-3, 4)))

    # real bid/ask of the file, 2012-03+ (report-only)
    raw = pd.read_parquet(B.CFD)
    raw.columns = [c.lower() for c in raw.columns]
    raw = raw.set_index("date")
    idx = pd.DatetimeIndex(raw.index).tz_localize("America/Chicago", ambiguous="NaT", nonexistent="NaT")
    raw = raw[~idx.isna()]
    raw.index = idx[~idx.isna()].tz_convert(NY) - pd.Timedelta(1, "min")
    raw = raw[~raw.index.duplicated(keep="last")]
    rm = raw.index.hour * 60 + raw.index.minute
    s_in = (raw["ask_open"] - raw["bid_open"])[rm == 930]
    s_out = (raw["ask_close"] - raw["bid_close"])[rm == 959]
    s_in.index, s_out.index = s_in.index.normalize().tz_localize(None), s_out.index.normalize().tz_localize(None)
    E = D[D.index >= "2012-03-01"].copy()
    E["real"] = E["R"] - 0.5 * (s_in.reindex(E.index) + s_out.reindex(E.index)) / E["risk"]
    print(row("report-only: real bid/ask 2012-03+", E, "real") + f" | spread in/out median {s_in.reindex(E.index).median():.2f}/{s_out.reindex(E.index).median():.2f} pt")

    print("\n2. PLACEBO AND LEAK CHECKS (dataset 1)")
    for lag in (1, 2):
        s = D["sig"].shift(lag)
        r = pd.Series(np.where(s > 0, D["R_long"], D["R_short"]), index=D.index)[s.notna()]
        print(f"  signal lagged {lag} day(s): gross mean {r.mean():+.4f}R  t={tstat(r):+.2f}")
    print(f"  always long: gross {D['R_long'].mean():+.4f}R t={tstat(D['R_long']):+.2f} | always short: {D['R_short'].mean():+.4f}R t={tstat(D['R_short']):+.2f}")
    RL, RS, n = D["R_long"].to_numpy(), D["R_short"].to_numpy(), len(D)
    obs, ge = D["R"].mean(), 0
    for _ in range(100):
        pick = RNG.integers(0, 2, size=(100, n)).astype(bool)
        ge += int((np.where(pick, RL, RS).mean(axis=1) >= obs).sum())
    p_rand = ge / 10000
    print(f"  random signs, 10,000 draws: p = {p_rand:.4f}")

    p1 = pd.read_csv(os.path.join(LAB, "_bar_boundary_trueclock_trades.csv"))
    p1["day"] = pd.to_datetime(p1["entry_ts"], utc=True).dt.tz_convert(NY).dt.tz_localize(None).dt.normalize()
    j = pd.concat([D["net"].rename("x"), p1.groupby("day")["net"].sum().rename("p")], axis=1)
    j = j[j["x"].notna()]
    rho_all, rho_both = j["x"].corr(j["p"].fillna(0.0)), j.dropna()["x"].corr(j.dropna()["p"])
    # [withheld: refers to the lab's private live strategy]

    print("\n3. REPLICATION A: LSE NQ.F M1 2016-06 -> 2026-10, futures cost")
    d = pd.read_parquet(LSE)
    d["timestamp"] = pd.to_datetime(d["timestamp"], utc=True)
    m = d.drop_duplicates("timestamp").set_index("timestamp").sort_index()[["open", "high", "low", "close"]].astype(float)
    m.index = m.index.tz_convert(NY)
    A = add_cost(build(m), FUT_PTS, ref)
    A.to_csv(os.path.join(OUT, "_idmom02_trades_lse.csv"))
    print(row("gross", A, "R"))
    print(row("net futures 1x", A))
    print(row("net futures 2x", A, "net2"))
    print(row("net at CFD cost", add_cost(A, CFD_PTS, ref)))
    print(row("2019-2026 net futures", A[A.year >= 2019]))
    c = pd.concat([D["R"].rename("cfd"), A["R"].rename("lse")], axis=1).dropna()
    same = (np.sign(D["rod"]).reindex(c.index) == np.sign(A["rod"]).reindex(c.index)).mean()
    corr_a = c["cfd"].corr(c["lse"])
    print(f"  common days {len(c)} | same signal {100 * same:.1f}% | corr of daily R with dataset 1: {corr_a:.3f}")
    print(row("dataset 1 on the same days, net CFD", D.loc[c.index]))
    print("  by year (net futures, sum R / t): " + "  ".join(f"{y}:{g['net'].sum():+.0f}/{tstat(g['net']):+.1f}" for y, g in A.groupby("year")))

    print("\n4. REPLICATION B and ES (Databento 1s -> M1, 2013-2018)")
    for sym, pts, rf in [("NQ", FUT_PTS, ref), ("ES", 0.25 + 2 * 2.25 / 50.0, 6700.0)]:
        m1, pc = dbn_m1(sym)
        X = add_cost(build(m1, prev_close=pc), pts, rf)
        X = X[X.year >= 2013]
        X.to_csv(os.path.join(OUT, f"_idmom02_trades_dbn_{sym}.csv"))
        print(row(f"{sym} gross", X, "R"))
        print(row(f"{sym} net futures 1x", X) + f" | cost median {X['cost'].median():.4f}R")
        c = pd.concat([D["R"].rename("cfd"), X["R"].rename("x")], axis=1).dropna()
        print(f"    common days with dataset 1: {len(c)} | corr of daily R {c['cfd'].corr(c['x']):.3f} | by year: " +
              "  ".join(f"{y}:{g['net'].sum():+.0f}" for y, g in X.groupby("year")))

    t1, t2 = tstat(D["net"]), tstat(D["net2"])
    ta, tb, tc = tstat(D[D.year <= 2015]["net"]), tstat(D[D.year >= 2016]["net"]), tstat(D[D.year >= 2019]["net"])
    s1 = D["sig"].shift(1)
    tl = tstat(pd.Series(np.where(s1 > 0, D["R_long"], D["R_short"]), index=D.index)[s1.notna()])
    k = {"K0": len(D) >= 2000, "K1": t1 >= 2.0, "K2": min(ta, tb, tc) >= 1.0, "K3": t2 >= 1.0,
         "K4": abs(tl) < 2 and p_rand < 0.01, "K5": abs(rho_all) < 0.20,
         "K6": tstat(A["net"]) >= 1.0 and A["net"].mean() > 0 and corr_a >= 0.90}
    print(f"\nGATES  K0 n={len(D)} | K1 t={t1:+.2f} | K2 {ta:+.2f}/{tb:+.2f}/{tc:+.2f} | K3 2x t={t2:+.2f} | K4 lag1 t={tl:+.2f}, p={p_rand:.4f} | "
          f"K5 rho={rho_all:+.3f} | K6 LSE t={tstat(A['net']):+.2f}, corr={corr_a:.3f}")
    print("       " + "  ".join(f"{a} {'PASS' if b else 'FAIL'}" for a, b in k.items()))
    print("VERDICT:", ("CANDIDATE" if t1 >= 3.43 else "FLAG (forward paper only)") if all(k.values()) else "KILL / REJECT")


if __name__ == "__main__":
    main()