#!/usr/bin/env python3
"""Three-layer overnight study: technical × event × behavioural.

Bins, thresholds and decision rules are frozen in
2026-09-16-three-layer-overnight-prereg.md. Nothing here may be re-tuned after
seeing output; a result that needs a different threshold is a fragility finding.

    .venv/bin/python studies/run_three_layer_overnight.py [--json out.json]
"""
from __future__ import annotations

import json
import random
import statistics
import sys
from collections import defaultdict
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
from zoneinfo import ZoneInfo

sys.path.insert(0, str(Path(__file__).resolve().parent))
from run_overnight_event_nights import (  # noqa: E402
    CTX, DATA, ET, QUADS, SPEC, load_all_events, load_weights, index_nights, parse_bar_ts,
)

import numpy as np  # noqa: E402

YEARS = range(2016, 2027)
SPLIT = "2022-01-01"           # train 2016-2021 | confirm 2022-2026
ATTN_Z = 1.0                   # frozen
SIGMA_MULT = 1.5               # frozen
VIX_CUTS = (15.0, 25.0)        # frozen
TRAIL = 20                     # frozen
CONSEC_DOWN = 2                # frozen
BONFERRONI_N = 15              # frozen: T1-T4, N1-N8 (minus GC-only), B1-B4
ALPHA = 0.05 / BONFERRONI_N
random.seed(20260916)
rng = np.random.default_rng(20260916)


# ── data ────────────────────────────────────────────────────────────────────
def load_sessions(root: str):
    """{date: {open, close}} for RTH, from the same bars the harness uses."""
    day: dict[str, dict] = {}
    for y in YEARS:
        p = DATA / "glbx" / root / "1m" / f"{y}.csv"
        if not p.exists():
            continue
        with p.open() as f:
            next(f)
            for line in f:
                if not line.strip():
                    continue
                c = line.split(",")
                try:
                    t = parse_bar_ts(c[0])
                    o = float(c[4]); cl = float(c[7])
                except (ValueError, IndexError):
                    continue
                mins = t.hour * 60 + t.minute
                if mins not in (9 * 60 + 30, 16 * 60, 15 * 60 + 59):
                    continue
                d = t.date().isoformat()
                rec = day.setdefault(d, {})
                if mins == 9 * 60 + 30:
                    rec["open"] = o
                else:
                    rec["close"] = cl
    return day


def load_vix():
    rows = json.loads((DATA / "fmp" / "vix-daily.json").read_text())
    return {r["date"]: float(r["close"]) for r in rows if r.get("close") is not None}


def load_features():
    out = {}
    p = CTX / "features" / "context_daily.jsonl"
    for line in p.read_text().splitlines():
        if line.strip():
            r = json.loads(line)
            out[r["date"]] = r
    return out


def build_nights(root: str, feats, vix):
    """One row per night: net $, technical state, event flags, behavioural flags."""
    day = load_sessions(root)
    dates = sorted(day)
    spec = SPEC[root]
    pv = spec["per_tick"] / spec["tick"]
    rth: dict[str, float] = {}
    for d in dates:
        r = day[d]
        if "open" in r and "close" in r:
            rth[d] = r["close"] - r["open"]
    rows = []
    for i in range(1, len(dates)):
        prev, cur = dates[i - 1], dates[i]
        if cur in QUADS or prev in QUADS:
            continue
        a, b = day[prev], day[cur]
        if "close" not in a or "open" not in b or prev not in rth:
            continue
        net = (b["open"] - a["close"]) * pv - spec["cost"]

        # technical state, all from sessions strictly BEFORE the night starts
        hist = [rth[d] for d in dates[max(0, i - 1 - TRAIL):i - 1] if d in rth]
        sd = statistics.stdev(hist) if len(hist) > 2 else float("nan")
        prev_ret = rth[prev]
        big = (not np.isnan(sd)) and abs(prev_ret) > SIGMA_MULT * sd
        consec = 0
        for d in reversed(dates[:i]):
            if d in rth and rth[d] < 0:
                consec += 1
            else:
                break
        v = vix.get(prev)
        vix_bucket = "na" if v is None else ("low" if v < VIX_CUTS[0] else "high" if v > VIX_CUTS[1] else "mid")
        weekend = datetime.fromisoformat(cur).weekday() == 0  # Monday morning = Friday night

        f = feats.get(cur, {})
        z = f.get("attn_nq_z")
        rows.append({
            "date": cur, "prev": prev, "net": net,
            # T
            "T1_down": prev_ret < 0,
            "T2_big": bool(big),
            "T3_vix": vix_bucket,
            "T4_weekend": weekend,
            # N (from the frozen event arms)
            "N1_prefomc": bool(f.get("prefomc")),
            "N2_amc_heavy": bool(f.get("amc_heavy")),
            "N3_amc_light": bool(f.get("amc_light")),
            "N4_fed_speech": bool(f.get("fed_speech_after16")),
            "N5_cpi_pce_nfp": bool(f.get("cpi_pce_nfp")),
            "N6_eia_wasde": bool(f.get("eia_wasde")),
            "N7_policy_geo": bool(f.get("policy_geo")),
            "N8_quiet": bool(f.get("quiet")),
            # B
            "B1_attn": (z is not None and z >= ATTN_Z),
            "B2_attn_down": (z is not None and z >= ATTN_Z and prev_ret < 0),
            "B3_chase": bool(big and prev_ret > 0),
            "B4_consec_down": consec >= CONSEC_DOWN,
            "attn_z": z,
        })
    return rows


# ── statistics ──────────────────────────────────────────────────────────────
def tstat(xs):
    if len(xs) < 3:
        return float("nan")
    sd = statistics.stdev(xs)
    return float("nan") if sd == 0 else statistics.mean(xs) / (sd / len(xs) ** 0.5)


def p_two_sided(t, n):
    """Normal approximation; n is large everywhere here."""
    if n < 3 or np.isnan(t):
        return float("nan")
    from math import erf, sqrt
    return 2 * (1 - 0.5 * (1 + erf(abs(t) / sqrt(2))))


def boot_ci(xs, reps=10000):
    if len(xs) < 5:
        return (float("nan"), float("nan"))
    a = np.array(xs, dtype=float)
    idx = rng.integers(0, len(a), size=(reps, len(a)))
    means = a[idx].mean(axis=1)
    return (float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)))


def placebo_p(xs, all_rows, reps=1000):
    """Share of equal-size random night samples whose mean beats this arm's."""
    if not xs:
        return float("nan")
    pool = np.array([r["net"] for r in all_rows], dtype=float)
    k, obs = len(xs), statistics.mean(xs)
    draws = rng.choice(pool, size=(reps, k), replace=True).mean(axis=1)
    return float((np.abs(draws - pool.mean()) >= abs(obs - pool.mean())).mean())


def arm(rows, pred, all_rows, label):
    xs = [r["net"] for r in rows if pred(r)]
    if not xs:
        return {"label": label, "n": 0}
    t = tstat(xs)
    lo, hi = boot_ci(xs)
    return {
        "label": label, "n": len(xs), "mean": statistics.mean(xs), "t": t,
        "p": p_two_sided(t, len(xs)),
        "win": 100.0 * sum(1 for x in xs if x > 0) / len(xs),
        "ci": [lo, hi], "placebo_p": placebo_p(xs, all_rows),
        "total": sum(xs),
    }


def split_signs(rows, pred):
    tr = [r["net"] for r in rows if pred(r) and r["date"] < SPLIT]
    cf = [r["net"] for r in rows if pred(r) and r["date"] >= SPLIT]
    m1 = statistics.mean(tr) if len(tr) > 2 else float("nan")
    m2 = statistics.mean(cf) if len(cf) > 2 else float("nan")
    agree = (not np.isnan(m1) and not np.isnan(m2) and (m1 > 0) == (m2 > 0))
    return {"train_n": len(tr), "train_mean": m1, "confirm_n": len(cf), "confirm_mean": m2, "sign_agree": bool(agree)}


def verdict(a, s):
    if a.get("n", 0) < 20 or np.isnan(a.get("t", float("nan"))):
        return "inconclusive (n too small)"
    passed = a["p"] < ALPHA and a["placebo_p"] < 0.05
    if passed and s["sign_agree"]:
        return "VALIDATED"
    if a["p"] < 0.05:
        return "suggestive (not tradeable)" if s["sign_agree"] else "refuted (sign flips between halves)"
    return "refuted"


def ols(rows):
    """Net ~ blocks, together. Robust (HC1) standard errors."""
    cols = ["T1_down", "T2_big", "T4_weekend", "N1_prefomc", "N2_amc_heavy", "N3_amc_light",
            "N4_fed_speech", "N5_cpi_pce_nfp", "N7_policy_geo", "B1_attn", "B3_chase", "B4_consec_down"]
    use = [r for r in rows if r["T3_vix"] != "na"]
    X = np.column_stack(
        [np.ones(len(use))]
        + [np.array([1.0 if r[c] else 0.0 for r in use]) for c in cols]
        + [np.array([1.0 if r["T3_vix"] == b else 0.0 for r in use]) for b in ("mid", "high")]
    )
    y = np.array([r["net"] for r in use])
    names = ["const"] + cols + ["T3_vix_mid", "T3_vix_high"]
    xtx_inv = np.linalg.pinv(X.T @ X)
    beta = xtx_inv @ X.T @ y
    resid = y - X @ beta
    n, k = X.shape
    meat = (X * (resid ** 2)[:, None]).T @ X * (n / max(n - k, 1))
    cov = xtx_inv @ meat @ xtx_inv
    se = np.sqrt(np.diag(cov))
    return [{"name": nm, "coef": float(b), "se": float(s), "t": float(b / s) if s else float("nan"),
             "p": p_two_sided(b / s if s else float("nan"), n)} for nm, b, s in zip(names, beta, se)], len(use)


def main():
    feats, vix = load_features(), load_vix()
    out = {"frozen": "2026-09-16-three-layer-overnight-prereg.md", "alpha_bonferroni": ALPHA, "roots": {}}
    for root in ("NQ", "ES", "GC"):
        rows = build_nights(root, feats, vix)
        if not rows:
            continue
        base = [r["net"] for r in rows]
        res = {"n_nights": len(rows), "baseline_mean": statistics.mean(base), "baseline_t": tstat(base),
               "baseline_total": sum(base), "first": rows[0]["date"], "last": rows[-1]["date"], "arms": []}
        tests = [
            ("T1 after a down RTH day", lambda r: r["T1_down"]),
            ("T1b after an up RTH day", lambda r: not r["T1_down"]),
            ("T2 after a >1.5sigma move", lambda r: r["T2_big"]),
            ("T3 VIX < 15", lambda r: r["T3_vix"] == "low"),
            ("T3 VIX 15-25", lambda r: r["T3_vix"] == "mid"),
            ("T3 VIX > 25", lambda r: r["T3_vix"] == "high"),
            ("T4 Friday->Monday night", lambda r: r["T4_weekend"]),
            ("N1 pre-FOMC", lambda r: r["N1_prefomc"]),
            ("N2 AMC earnings heavy", lambda r: r["N2_amc_heavy"]),
            ("N3 AMC earnings light", lambda r: r["N3_amc_light"]),
            ("N4 Fed speech after 16:00", lambda r: r["N4_fed_speech"]),
            ("N5 CPI/PCE/NFP tomorrow", lambda r: r["N5_cpi_pce_nfp"]),
            ("N6 EIA/WASDE tomorrow", lambda r: r["N6_eia_wasde"]),
            ("N7 policy/geo headline", lambda r: r["N7_policy_geo"]),
            ("N8 quiet night", lambda r: r["N8_quiet"]),
            ("B1 attention z >= 1.0", lambda r: r["B1_attn"]),
            ("B2 attention x down day", lambda r: r["B2_attn_down"]),
            ("B3 chase (big up day)", lambda r: r["B3_chase"]),
            ("B4 2+ down days", lambda r: r["B4_consec_down"]),
        ]
        for label, pred in tests:
            a = arm(rows, pred, rows, label)
            if a.get("n", 0):
                s = split_signs(rows, pred)
                a.update(s)
                a["verdict"] = verdict(a, s)
            res["arms"].append(a)
        if root == "NQ":
            # placebo (b): shuffle the attention series across dates, re-run B1
            zs = [r["attn_z"] for r in rows]
            shuffled = list(zs)
            random.shuffle(shuffled)
            fake = [dict(r, attn_z=z, B1_attn=(z is not None and z >= ATTN_Z)) for r, z in zip(rows, shuffled)]
            res["placebo_shuffled_attention"] = arm(fake, lambda r: r["B1_attn"], rows, "B1 with dates shuffled")
            coefs, n_used = ols(rows)
            res["ols"] = {"n": n_used, "coefs": coefs}
        out["roots"][root] = res

    if "--json" in sys.argv:
        Path(sys.argv[sys.argv.index("--json") + 1]).write_text(json.dumps(out, indent=1))
    for root, res in out["roots"].items():
        print(f"\n=== {root}  n={res['n_nights']} nights {res['first']}..{res['last']}  "
              f"baseline ${res['baseline_mean']:.2f}/night  t={res['baseline_t']:.2f}  total ${res['baseline_total']:,.0f}")
        print(f"{'arm':34s} {'n':>5s} {'mean$':>9s} {'t':>6s} {'p':>8s} {'win%':>6s} {'plc':>6s}  split  verdict")
        for a in res["arms"]:
            if not a.get("n"):
                continue
            print(f"{a['label']:34s} {a['n']:5d} {a['mean']:9.2f} {a['t']:6.2f} {a['p']:8.4f} "
                  f"{a['win']:6.1f} {a['placebo_p']:6.3f}  {'same' if a['sign_agree'] else 'FLIP'}  {a['verdict']}")
        if "ols" in res:
            print(f"\n  OLS (n={res['ols']['n']}, HC1):")
            for c in res["ols"]["coefs"]:
                star = " *" if c["p"] < ALPHA else ""
                print(f"    {c['name']:18s} {c['coef']:9.2f} ± {c['se']:7.2f}  t={c['t']:6.2f}  p={c['p']:.4f}{star}")
            p = res["placebo_shuffled_attention"]
            print(f"\n  placebo, attention shuffled: n={p['n']} mean=${p['mean']:.2f} t={p['t']:.2f} p={p['p']:.4f}")
    print(f"\nBonferroni threshold p < {ALPHA:.4f} (15 frozen hypotheses)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())