# =============================================================================
#  Project   : anomaly-atlas
#  File      : experiments/micro/expF_multiple_testing/benchmark.py
#  Purpose   : Survival battery: naive -> FDR -> RC/SPA -> DSR over C/D/E rules
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Data src  : hfmarketdata.io (sole data source)
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Experiment F — the survival curve (protocol pre-specified in hypothesis.md).

Rules are built mechanically from EVERYTHING the C/D/E scans searched (both
signs), evaluated on the same TRAIN data (in-sample by design — OOS is expH),
and pushed through naive-t -> BH-FDR -> White RC / Hansen SPA -> DSR.
"""

from __future__ import annotations

import json
import sys
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path

import numpy as np
from scipy.stats import norm

REPO_ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO_ROOT / "benchmarks"))
sys.path.insert(0, str(REPO_ROOT / "src"))

from hardware_manifest import collect_manifest  # noqa: E402

from anomaly_atlas.data.cleaning import RTH_SLOTS, rth_day_grids  # noqa: E402
from anomaly_atlas.data.hf_client import HFMarketDataClient  # noqa: E402
from anomaly_atlas.data.universe import TRAIN, TRAIN_SUBPERIODS, core_universe  # noqa: E402
from anomaly_atlas.stats.multiple_testing import benjamini_hochberg  # noqa: E402
from anomaly_atlas.stats.spa import deflated_sharpe, reality_check, spa_test  # noqa: E402

ADJ = "adj_split"
N_BOOT, MEAN_BLOCK, SEED = 500, 5.0, 42
LIQUID = ["AAPL", "MSFT", "NVDA", "AMZN", "GOOGL", "META", "TSLA", "JPM", "XOM", "UNH", "QQQ"]
SECTORS = ["XLF", "XLE", "XLK", "XLV", "XLI", "XLY", "XLP", "XLU", "XLB"]
D_WINDOWS = {"2006-2007": ("2006-01-01", "2008-01-01"),
             "2014-2015": ("2014-01-01", "2016-01-01")}
ONE_MIN_WINDOW = ("2014-01-01", "2016-01-01")


def contrarian_daily(bars: list[dict], timeframe: str) -> dict[str, float]:
    """day -> contrarian rule return at the cell's timeframe (trade-time RTH)."""
    by_day: dict[str, list[float]] = defaultdict(list)
    for b in bars:
        dt = b["datetime"]
        if timeframe == "1day":
            by_day[dt[:10]].append(np.log(b["close"]))
        elif "09:30" <= dt[11:16] < "16:00":
            by_day[dt[:10]].append(np.log(b["close"]))
    days = sorted(by_day)
    out: dict[str, float] = {}
    if timeframe == "1day":
        closes = np.array([by_day[d][0] for d in days])
        r = np.diff(closes)
        for i in range(1, len(r)):
            out[days[i + 1]] = float(-np.sign(r[i - 1]) * r[i])
        return out
    for d in days:
        p = np.array(by_day[d])
        if len(p) < 3:
            continue
        r = np.diff(p)
        out[d] = float(np.sum(-np.sign(r[:-1]) * r[1:]))
    return out


def leadlag_daily(gx: dict, gy: dict, day_list: list[str]) -> dict[str, float]:
    """day -> sign(leader_{t-1}) * follower_t summed over both-fresh minutes."""
    out: dict[str, float] = {}
    for day in day_list:
        px, py = gx.get(day), gy.get(day)
        if px is None or py is None:
            continue
        ox, oy = np.isfinite(px), np.isfinite(py)
        fx = np.zeros(RTH_SLOTS - 1, dtype=bool)
        fy = np.zeros(RTH_SLOTS - 1, dtype=bool)
        rx = np.zeros(RTH_SLOTS - 1)
        ry = np.zeros(RTH_SLOTS - 1)
        fpx, fpy = px.copy(), py.copy()
        for t in range(1, RTH_SLOTS):
            if not ox[t]:
                fpx[t] = fpx[t - 1]
            if not oy[t]:
                fpy[t] = fpy[t - 1]
        rx[:] = np.diff(fpx)
        ry[:] = np.diff(fpy)
        fx[:] = ox[1:] & ox[:-1]
        fy[:] = oy[1:] & oy[:-1]
        keep = fx[:-1] & fy[1:] & np.isfinite(rx[:-1]) & np.isfinite(ry[1:])
        if keep.sum() < 30:
            continue
        out[day] = float(np.sum(np.sign(rx[:-1][keep]) * ry[1:][keep]))
    return out


def build_matrix(rules: dict[str, dict[str, float]]) -> tuple[np.ndarray, list[str], list[str]]:
    """(days x 2N signed rules) matrix; missing day = 0 (idle), as declared."""
    day_set = sorted({d for r in rules.values() for d in r})
    names, cols = [], []
    for name, series in rules.items():
        v = np.array([series.get(d, 0.0) for d in day_set])
        names += [f"+{name}", f"-{name}"]
        cols += [v, -v]
    return np.column_stack(cols), names, day_set


def battery(x: np.ndarray, names: list[str]) -> dict:
    t_len = x.shape[0]
    mu = x.mean(axis=0)
    sd = x.std(axis=0, ddof=1)
    sd = np.maximum(sd, 1e-12)
    t = np.sqrt(t_len) * mu / sd
    p_two = 2 * (1 - norm.cdf(np.abs(t)))
    rc = reality_check(x, n_boot=N_BOOT, mean_block=MEAN_BLOCK, seed=SEED)
    sp = spa_test(x, n_boot=N_BOOT, mean_block=MEAN_BLOCK, seed=SEED)
    srs = mu / sd
    best = int(np.argmax(srs))
    r = x[:, best]
    dsr = deflated_sharpe(
        sr=float(srs[best]), t_len=t_len,
        skew=float(((r - r.mean()) ** 3).mean() / r.std() ** 3),
        kurt=float(((r - r.mean()) ** 4).mean() / r.std() ** 4),
        n_trials=x.shape[1], sr_variance=float(srs.var(ddof=1)),
    )
    return {
        "days": int(t_len), "rules": len(names),
        "naive_t196": int((np.abs(t) > 1.96).sum()),
        "p_two_sided": p_two, "t": t,
        "rc_p": rc["p"], "spa_p": sp["p"],
        "spa_step1_survivors": [names[i] for i, tv in enumerate(sp["rule_t"])
                                if tv >= sp["t95"]],
        "best_rule": names[best], "best_daily_sharpe": round(float(srs[best]), 4),
        "best_annualized_sharpe": round(float(srs[best] * np.sqrt(252)), 3),
        "dsr_sr0": round(dsr["sr0"], 4), "dsr": round(dsr["dsr"], 4),
    }


def main() -> None:
    run_utc = datetime.now(UTC)
    client = HFMarketDataClient()
    universe = core_universe(client.tickers("stock", timeframe="1min", adjustment=ADJ))

    blocks: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)

    # ---- R-family (expC universe)
    for asset, ticker, _ in universe:
        for sub, (s, e) in TRAIN_SUBPERIODS.items():
            day_bars = client.get_bars(asset, ticker, "1day", ADJ, s, e)
            if len(day_bars) < 200:
                continue
            for tf in ("1day", "30min", "5min"):
                bars = day_bars if tf == "1day" else client.get_bars(asset, ticker, tf, ADJ, s, e)
                series = contrarian_daily(bars, tf)
                if len(series) >= 150:
                    blocks[f"expC {sub}"][f"R:{ticker}:{tf}"] = series
        print(f"R {ticker} done")
    for asset, ticker in [("stock", t) for t in LIQUID if t != "QQQ"] + \
                         [("etf", t) for t in ("SPY", "QQQ")]:
        bars = client.get_bars(asset, ticker, "1min", ADJ, *ONE_MIN_WINDOW)
        series = contrarian_daily(bars, "1min")
        if len(series) >= 150:
            blocks["expC 1min 2014-2015"][f"R:{ticker}:1min"] = series

    # ---- L-family (expD universe)
    random10 = sorted(np.random.default_rng(42).choice(
        sorted(client.tickers("stock", timeframe="1min", adjustment=ADJ)), 30,
        replace=False))[:10]
    for window, (s, e) in D_WINDOWS.items():
        spec = ({"SPY": ("etf", "SPY", ADJ)}
                | {t: ("stock", t, ADJ) for t in LIQUID if t != "QQQ"}
                | {"QQQ": ("etf", "QQQ", ADJ)}
                | {t: ("etf", t, ADJ) for t in SECTORS}
                | {t: ("stock", t, ADJ) for t in random10})
        if window == "2014-2015":
            spec |= {"SPX": ("index", "SPX", None),
                     "ES": ("futures", "ES", "contin_adj_ratio")}
        grids = {}
        for name, (asset, ticker, adj) in spec.items():
            g = rth_day_grids(client.get_bars(asset, ticker, "1min", adj, s, e))
            if len(g) >= 200:
                grids[name] = g
        day_list = sorted(grids["SPY"].keys())
        for name, g in grids.items():
            if name == "SPY":
                continue
            if name in ("SPX", "ES"):
                series = leadlag_daily(g, grids["SPY"], day_list)  # x leads SPY
                key = f"L:{name}->SPY"
            else:
                series = leadlag_daily(grids["SPY"], g, day_list)
                key = f"L:SPY->{name}"
            if len(series) >= 150:
                blocks[f"expD {window}"][key] = series
        print(f"L {window} done ({len(blocks[f'expD {window}'])} pairs)")

    # ---- C-family (expE classes, drift-adjusted)
    bars = client.get_bars("etf", "SPY", "1day", "adj_splitdiv", TRAIN[0], TRAIN[1])
    dates = [b["datetime"][:10] for b in bars][1:]
    rets = np.diff(np.log([b["close"] for b in bars]))
    mu = rets.mean()
    sys.path.insert(0, str(REPO_ROOT / "experiments" / "micro" / "expE_calendar_scan"))
    from benchmark import class_masks  # noqa: E402

    masks = class_masks(dates)
    for cname, m in masks.items():
        blocks["expE train"][f"C:{cname}"] = {
            d: float(rets[i] - mu) for i, d in enumerate(dates) if m[i]
        }

    # ---- battery per block + global funnel
    per_block: dict[str, dict] = {}
    all_p, all_index = [], []
    for bname, rules in blocks.items():
        x, names, _ = build_matrix(rules)
        res = battery(x, names)
        all_p.extend(res.pop("p_two_sided").tolist())
        res.pop("t")
        all_index.extend([(bname, n) for n in names])
        per_block[bname] = res
        print(f"{bname}: rules={res['rules']} naive={res['naive_t196']} "
              f"rc_p={res['rc_p']} spa_p={res['spa_p']} dsr={res['dsr']}")

    fdr_mask = benjamini_hochberg(np.array(all_p), alpha=0.05)
    total_rules = len(all_p)
    funnel = {
        "universe_rules": total_rules,
        "naive_t196": int(sum(per_block[b]["naive_t196"] for b in per_block)),
        "fdr_survivors": int(fdr_mask.sum()),
        "spa_step1_survivors": sorted({n for b in per_block
                                       for n in per_block[b]["spa_step1_survivors"]}),
        "blocks_spa_significant": [b for b in per_block if per_block[b]["spa_p"] < 0.05],
        "dsr_by_block": {b: per_block[b]["dsr"] for b in per_block},
    }
    funnel["survival_rate"] = {
        "naive": round(funnel["naive_t196"] / total_rules, 4),
        "fdr": round(funnel["fdr_survivors"] / total_rules, 4),
        "spa_step1": round(len(funnel["spa_step1_survivors"]) / total_rules, 4),
    }

    results = {
        "experiment": "expF_multiple_testing",
        "run_utc": run_utc.isoformat(),
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "data_source": "hfmarketdata.io",
        "confidence_level": 0,
        "protocol": {"n_boot": N_BOOT, "mean_block": MEAN_BLOCK, "seed": SEED,
                     "note": "in-sample search survival on TRAIN; OOS = expH"},
        "per_block": per_block,
        "funnel": funnel,
        "client_stats": vars(client.stats) | {"refreshes": list(client.stats.refreshes)},
        "manifest": collect_manifest(),
    }
    out_dir = REPO_ROOT / "results" / "expF_multiple_testing" / run_utc.strftime("%Y%m%dT%H%M%SZ")
    out_dir.mkdir(parents=True)
    (out_dir / "results.json").write_text(json.dumps(results, indent=2) + "\n")
    print(f"\nwrote {out_dir.relative_to(REPO_ROOT)}/results.json")
    print(json.dumps(funnel, indent=1))


if __name__ == "__main__":
    main()