#!/usr/bin/env python
"""Macro events, surprises and administrations vs NQ.

Policy frozen BEFORE results:
  <data-root>/fp-data/studies/2026-08-25-fundamentals-prereg.md

Reuses the shipped desk FAMILIES/econSign table from
fp-terminal/scripts/backtest-surprise.mjs verbatim so numbers stay comparable,
but standardizes surprises with an EXPANDING (strictly backward) sigma rather
than the full-sample sigma the shipped script uses.
"""
import glob
import json
import re
from datetime import date, datetime, timezone

import numpy as np
import pandas as pd

import rvol_orb as R

X10 = "<data-root>/fp-data"
ET = "America/New_York"

# verbatim from fp-terminal/scripts/backtest-surprise.mjs
FAMILIES = [
    ("core_cpi_mom", r"^(core cpi|core inflation rate) mom", -1),
    ("cpi_mom",      r"^(cpi|inflation rate) mom",           -1),
    ("core_pce_mom", r"^core pce price index mom",           -1),
    ("pce_mom",      r"^pce price index mom",                -1),
    ("core_ppi_mom", r"^core producer price index mom",      -1),
    ("ppi_mom",      r"^producer price index mom",           -1),
    ("retail_mom",   r"^retail sales mom",                   +1),
    ("durable_mom",  r"^durable goods orders mom",           +1),
    ("claims",       r"^initial jobless claims",             -1),
    ("nfp",          r"non.?farm payrolls",                   0),
    ("unemp",        r"^unemployment rate",                   0),
]
# tier-1 = the releases that reliably move index futures
TIER1 = r"^(cpi|inflation rate) mom|^core (cpi|inflation rate) mom|non.?farm payrolls|" \
        r"^unemployment rate|^core pce price index mom|^ism (manufacturing|services) pmi|" \
        r"^retail sales mom|^fed interest rate decision|^gross domestic product qoq"

ADMINS = [
    ("Obama-II", date(2010, 1, 1),  date(2017, 1, 20)),
    ("Trump-I",  date(2017, 1, 20), date(2021, 1, 20)),
    ("Biden",    date(2021, 1, 20), date(2025, 1, 20)),
    ("Trump-II", date(2025, 1, 20), date(2100, 1, 1)),
]
MIN_FAM_OBS = 20
COST_PTS = 1.0


def admin_of(d):
    for name, lo, hi in ADMINS:
        if lo <= d < hi:
            return name
    return "?"


def strip_paren(ev):
    return re.sub(r"\s*\([^)]*\)\s*$", "", ev or "").strip()


def load_calendar():
    rows = []
    for f in sorted(glob.glob(f"{X10}/fmp/econ-calendar-*.json")):
        for r in json.load(open(f)):
            ev = strip_paren(r.get("event", ""))
            ts = r.get("date")
            if not ts:
                continue
            try:
                dt = datetime.strptime(ts, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
            except ValueError:
                continue
            et = dt.astimezone(pd.Timestamp(dt).tz_convert(ET).tz if False else None)
            et = pd.Timestamp(dt).tz_convert(ET)
            rows.append({
                "et": et, "date": et.date(), "tod": et.hour * 60 + et.minute,
                "event": ev, "impact": r.get("impact"),
                "country": str(r.get("country", "")).upper(),
                "estimate": pd.to_numeric(r.get("estimate"), errors="coerce"),
                "actual": pd.to_numeric(r.get("actual"), errors="coerce"),
            })
    df = pd.DataFrame(rows)
    return df[df["country"].isin(["US", "USA", "UNITED STATES"])].reset_index(drop=True)


def classify(ev):
    for key, pat, sign in FAMILIES:
        if re.search(pat, ev, re.I):
            return key, sign
    return None, None


def build_price():
    """NQ per-session targets."""
    sessions, meta = R.build_sessions("NQ")
    df = R.load_sym("NQ")
    out = []
    prev_18 = {}
    # prior-day 18:00 ET reference for the overnight/pre-open leg
    for d, day in df.groupby("date", sort=True):
        n = day[(day["tod"] >= 1080) & (day["tod"] < 1140)]   # 18:00-18:59 ET
        if len(n):
            prev_18[d] = float(n.iloc[0]["open"])
    dates = sorted(prev_18)
    prev_map = {dates[i]: dates[i - 1] for i in range(1, len(dates))}
    for s in sessions:
        d = s["date"]
        tape = meta["tapes"].get(d)
        day = df[df["date"] == d]
        rth = day[(day["tod"] >= 570) & (day["tod"] < 960)]
        if len(rth) < 330:
            continue
        o = float(rth.iloc[0]["open"])
        c = float(rth.iloc[-1]["close"])
        def at(tod):
            sel = rth[rth["tod"] == tod]
            return float(sel.iloc[0]["close"]) if len(sel) else np.nan
        pd_ = prev_map.get(d)
        pre = prev_18.get(pd_) if pd_ else None
        out.append({
            "date": d, "year": d.year, "admin": admin_of(d),
            "open": o, "close": c,
            "ret_rth": (c / o - 1) * 1e4,
            "ret_5m": (at(574) / o - 1) * 1e4 if not np.isnan(at(574)) else np.nan,
            "ret_15m": (at(584) / o - 1) * 1e4 if not np.isnan(at(584)) else np.nan,
            "ret_preopen": ((o / pre - 1) * 1e4) if pre else np.nan,
            "ret_c2c": np.nan,
        })
    p = pd.DataFrame(out).sort_values("date").reset_index(drop=True)
    p["ret_c2c"] = (p["close"] / p["close"].shift(1) - 1) * 1e4
    return p


def expanding_z(vals):
    """Strictly backward standardization; NaN until MIN_FAM_OBS priors exist."""
    out = np.full(len(vals), np.nan)
    for i in range(len(vals)):
        hist = vals[:i]
        hist = hist[~np.isnan(hist)]
        if len(hist) >= MIN_FAM_OBS:
            sd = hist.std(ddof=1)
            if sd > 0:
                out[i] = vals[i] / sd
    return out


def build_surprise(cal):
    cal = cal.dropna(subset=["estimate", "actual"]).copy()
    cal["surprise"] = cal["actual"] - cal["estimate"]
    cal[["fam", "sign"]] = cal["event"].apply(lambda e: pd.Series(classify(e)))
    cal = cal.dropna(subset=["fam"])
    cal = cal[cal["tod"] < 570]                       # pre-open only
    cal = cal.sort_values("et").drop_duplicates(["date", "fam"], keep="first")
    parts = []
    for fam, g in cal.groupby("fam"):
        g = g.sort_values("et").copy()
        g["z"] = expanding_z(g["surprise"].to_numpy(dtype=float))
        g["z_full"] = g["surprise"] / g["surprise"].std(ddof=1)
        parts.append(g)
    cal = pd.concat(parts).sort_values("et")
    comp = (cal[cal["sign"] != 0]
            .assign(sz=lambda x: x["sign"] * x["z"], szf=lambda x: x["sign"] * x["z_full"])
            .groupby("date")[["sz", "szf"]].sum().reset_index())
    sd = comp["sz"].expanding().std()
    comp["signed_z"] = comp["sz"] / sd
    comp["signed_z_full"] = comp["szf"] / comp["szf"].std(ddof=1)
    return cal, comp


def ols_t(x, y):
    m = ~(np.isnan(x) | np.isnan(y))
    x, y = x[m], y[m]
    n = len(x)
    if n < 10:
        return np.nan, np.nan, n
    b1 = np.cov(x, y, ddof=1)[0, 1] / np.var(x, ddof=1)
    b0 = y.mean() - b1 * x.mean()
    resid = y - (b0 + b1 * x)
    sxx = ((x - x.mean()) ** 2).sum()
    if sxx <= 0 or not np.isfinite(sxx):
        return b1, np.nan, n
    se = np.sqrt((resid @ resid / (n - 2)) / sxx)
    return (b1, b1 / se, n) if se > 0 else (b1, np.nan, n)


def tmean(a):
    a = np.asarray(a, float); a = a[~np.isnan(a)]
    if len(a) < 2:
        return np.nan
    return a.mean() / (a.std(ddof=1) / np.sqrt(len(a)))


def main():
    print("Loading calendar and NQ price history...")
    cal = load_calendar()
    px = build_price()
    print(f"calendar rows: {len(cal)}   NQ sessions: {len(px)}  "
          f"({px['date'].min()} -> {px['date'].max()})")

    ev, comp = build_surprise(cal)
    print(f"classified pre-open events w/ est+act: {len(ev)} across {ev['date'].nunique()} dates")

    m = px.merge(comp, on="date", how="left")
    print(f"sessions with a composite surprise: {m['signed_z'].notna().sum()}")

    print("\n" + "=" * 90)
    print("PRIMARY — OLS ret_rth ~ signed_z (pre-open composite), 2016-2026, bar |t| > 3.0")
    print("=" * 90)
    x = m["signed_z"].to_numpy(float); y = m["ret_rth"].to_numpy(float)
    b, t, n = ols_t(x, y)
    print(f"  n = {n}   slope = {b:+.2f} bps per 1sd surprise   t = {t:+.2f}")
    c1 = abs(t) > 3.0
    print(f"  >>> PRIMARY {'PASS' if c1 else 'FAIL'}")
    for tgt in ["ret_5m", "ret_15m", "ret_rth"]:
        b2, t2, n2 = ols_t(x, m[tgt].to_numpy(float))
        print(f"    {tgt:>10}: slope {b2:+7.2f} bps  t {t2:+6.2f}  n {n2}")
    bf, tf, nf = ols_t(m["signed_z_full"].to_numpy(float), y)
    print(f"  robustness (full-sample sigma, as shipped script): slope {bf:+.2f} t {tf:+.2f}")

    print("\n" + "=" * 90)
    print("SECONDARY 4 — is it tradable? sign(signed_z) traded at the open, exit 15:59")
    print("=" * 90)
    sub = m.dropna(subset=["signed_z", "ret_rth"]).copy()
    sub["dir"] = np.sign(sub["signed_z"])
    sub = sub[sub["dir"] != 0]
    pnl_bps = sub["dir"] * sub["ret_rth"]
    cost_bps = COST_PTS / sub["open"] * 1e4
    net = pnl_bps - cost_bps
    print(f"  n={len(sub)}  hit={100*(pnl_bps>0).mean():.1f}%  gross={pnl_bps.mean():+.2f} bps "
          f"t={tmean(pnl_bps):+.2f}")
    print(f"  net of {COST_PTS} pt: {net.mean():+.2f} bps  t={tmean(net):+.2f}  "
          f"(cost avg {cost_bps.mean():.2f} bps)")

    print("\n" + "=" * 90)
    print("SECONDARY 1 — event-day effect (Savor-Wilson analogue), 2016-2026")
    print("=" * 90)
    t1 = cal[cal["event"].str.contains(TIER1, case=False, regex=True, na=False)]  # noqa: match-groups ok
    t1d = set(t1["date"])
    px2 = px[px["date"] >= date(2016, 1, 1)].copy()
    px2["is_ev"] = px2["date"].isin(t1d)
    for tgt in ["ret_rth", "ret_c2c"]:
        a = px2[px2["is_ev"]][tgt].dropna(); b_ = px2[~px2["is_ev"]][tgt].dropna()
        diff = a.mean() - b_.mean()
        se = np.sqrt(a.var(ddof=1)/len(a) + b_.var(ddof=1)/len(b_))
        print(f"  {tgt:>8}: event n={len(a):4d} mean={a.mean():+7.2f} | "
              f"non-event n={len(b_):4d} mean={b_.mean():+7.2f} | diff={diff:+7.2f} t={diff/se:+.2f}")
    print(f"  (event days = {len(t1d & set(px2['date']))} of {len(px2)} sessions)")

    print("\n" + "=" * 90)
    print("SECONDARY 2 — pre-announcement drift: overnight into a tier-1 release")
    print("=" * 90)
    px2 = px2.copy()
    px2["ev_tomorrow"] = px2["is_ev"]
    a = px2[px2["is_ev"]]["ret_preopen"].dropna(); b_ = px2[~px2["is_ev"]]["ret_preopen"].dropna()
    se = np.sqrt(a.var(ddof=1)/len(a) + b_.var(ddof=1)/len(b_))
    print(f"  ret_preopen on event days n={len(a):4d} mean={a.mean():+6.2f} bps t={tmean(a):+.2f}")
    print(f"  ret_preopen other days   n={len(b_):4d} mean={b_.mean():+6.2f} bps t={tmean(b_):+.2f}")
    print(f"  difference {a.mean()-b_.mean():+.2f} bps  t={(a.mean()-b_.mean())/se:+.2f}")
    fomc = set(cal[cal["event"].str.contains("fed interest rate decision", case=False, na=False)]["date"])
    f = px2[px2["date"].isin(fomc)]["ret_preopen"].dropna()
    print(f"  FOMC-day pre-open (Lucca-Moench window analogue): n={len(f)} mean={f.mean():+.2f} t={tmean(f):+.2f}")

    print("\n" + "=" * 90)
    print("SECONDARY 3 — per-family slope on ret_rth (which single release moves the session)")
    print("=" * 90)
    print(f"{'family':>14} {'n':>5} {'slope':>9} {'t':>7}")
    for fam, _, sign in FAMILIES:
        g = ev[ev["fam"] == fam][["date", "z"]].dropna()
        j = px.merge(g, on="date", how="inner")
        if len(j) < 20:
            continue
        b3, t3, n3 = ols_t(j["z"].to_numpy(float), j["ret_rth"].to_numpy(float))
        star = " *" if abs(t3) > 3 else ""
        print(f"{fam:>14} {n3:>5} {b3:>+9.2f} {t3:>+7.2f}{star}   econSign={sign:+d}")

    print("\n" + "=" * 90)
    print("SECONDARY 5 — ADMINISTRATION (descriptive only; n=4 regimes, fully confounded)")
    print("=" * 90)
    print(f"{'admin':>10} {'sessions':>9} {'meanRTH':>9} {'t':>7} {'evDayRTH':>9} {'slope':>8} {'t':>7} {'nSurp':>6}")
    for name, lo, hi in ADMINS:
        sl = px[(px["date"] >= lo) & (px["date"] < hi)]
        if len(sl) == 0:
            continue
        ms = m[(m["date"] >= lo) & (m["date"] < hi)]
        bx, tx, nx = ols_t(ms["signed_z"].to_numpy(float), ms["ret_rth"].to_numpy(float))
        evm = px2[(px2["date"] >= lo) & (px2["date"] < hi) & px2["is_ev"]]["ret_rth"]
        print(f"{name:>10} {len(sl):>9} {sl['ret_rth'].mean():>+9.2f} {tmean(sl['ret_rth']):>+7.2f} "
              f"{(evm.mean() if len(evm) else np.nan):>+9.2f} {bx:>+8.2f} {tx:>+7.2f} {nx:>6}")
    print("  NOTE: administration is collinear with the Fed cycle, inflation regime, COVID and the")
    print("        mega-cap AI trend. These are era descriptions, not causal estimates.")

    print("\n" + "=" * 90)
    print("SECONDARY 6 — stability of the primary")
    print("=" * 90)
    for lab, sl in [("IS  (<2020)", m[m["year"] < 2020]), ("OOS (>=2020)", m[m["year"] >= 2020])]:
        b4, t4, n4 = ols_t(sl["signed_z"].to_numpy(float), sl["ret_rth"].to_numpy(float))
        print(f"  {lab}: n={n4:5d} slope={b4:+7.2f} t={t4:+6.2f}")
    pos = 0; yrs = sorted(m["year"].unique())
    for y in yrs:
        sl = m[m["year"] == y]
        b5, t5, n5 = ols_t(sl["signed_z"].to_numpy(float), sl["ret_rth"].to_numpy(float))
        if not np.isnan(b5):
            pos += b5 > 0
            print(f"    {y}: n={n5:4d} slope={b5:+8.2f} t={t5:+6.2f}")
    print(f"  years with positive slope: {pos}/{len([y for y in yrs if y>=2016])}")

    print("\n" + "=" * 90)
    print("VERDICT")
    print("=" * 90)
    print(f"  H1 primary |t|>3.0 ............ {'PASS' if c1 else 'FAIL'}  (t={t:+.2f})")
    print(f"  H1 tradable after cost ....... {'PASS' if tmean(net) > 3 else 'FAIL'}  (t={tmean(net):+.2f})")
    print(f"  H1 = {'CONFIRMED' if c1 and tmean(net) > 3 else 'REFUTED'}")


if __name__ == "__main__":
    main()