# =============================================================================
#  Project   : anomaly-atlas
#  File      : experiments/micro/expD_leadlag_scan/benchmark.py
#  Purpose   : Lead-lag scan — raw LOCF vs both-fresh, artifact share, FDR
#  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 D — lead-lag scan (protocol pre-specified in hypothesis.md).

The design measures T3 directly: every pair reports its raw-LOCF lead, its
both-fresh lead, and the difference (artifact share). Level 0 throughout.
"""

from __future__ import annotations

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

import numpy as np

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 (  # noqa: E402
    RTH_SLOTS,
    both_fresh,
    nan_xcorr,
    rth_day_grids,
)
from anomaly_atlas.data.hf_client import HFMarketDataClient  # noqa: E402
from anomaly_atlas.data.universe import random_stock_universe  # noqa: E402
from anomaly_atlas.stats.multiple_testing import (  # noqa: E402
    benjamini_hochberg,
    bootstrap_pvalue,
)

ADJ = "adj_split"
LIQUID = ["AAPL", "MSFT", "NVDA", "AMZN", "GOOGL", "META", "TSLA", "JPM", "XOM", "UNH", "QQQ"]
SECTORS = ["XLF", "XLE", "XLK", "XLV", "XLI", "XLY", "XLP", "XLU", "XLB"]
WINDOWS = {"2006-2007": ("2006-01-01", "2008-01-01"),
           "2014-2015": ("2014-01-01", "2016-01-01")}
MAX_LAG = 3
BOOT_N, BOOT_SEED = 200, 42
MIN_DAYS, MIN_FRESH = 200, 10_000


def day_matrix(days: dict[str, np.ndarray], day_list: list[str]) -> tuple[np.ndarray, np.ndarray]:
    """(n_days, 389) LOCF returns + interval-fresh masks, day-aligned."""
    rets = np.full((len(day_list), RTH_SLOTS - 1), np.nan)
    fresh = np.zeros((len(day_list), RTH_SLOTS - 1), dtype=bool)
    for i, day in enumerate(day_list):
        p = days.get(day)
        if p is None:
            continue
        observed = np.isfinite(p)
        filled = p.copy()
        for t in range(1, RTH_SLOTS):
            if not observed[t]:
                filled[t] = filled[t - 1]
        rets[i] = np.diff(filled)
        fresh[i] = observed[1:] & observed[:-1]
    return rets, fresh


def day_block_boot_xcorr(
    rx: np.ndarray, fx: np.ndarray, ry: np.ndarray, fy: np.ndarray, lag: int
) -> np.ndarray:
    """Bootstrap distribution of the both-fresh xcorr at `lag`, resampling days."""
    rng = np.random.default_rng(BOOT_SEED)
    n_days = rx.shape[0]
    out = np.empty(BOOT_N)
    for b in range(BOOT_N):
        idx = rng.integers(0, n_days, n_days)
        bx, by = both_fresh(rx[idx].ravel(), fx[idx].ravel(), ry[idx].ravel(), fy[idx].ravel())
        out[b] = nan_xcorr(bx, by, abs(lag)).get(lag, np.nan)
    return out


def analyze_pair(rx, fx, ry, fy) -> dict | None:
    bx, by = both_fresh(rx.ravel(), fx.ravel(), ry.ravel(), fy.ravel())
    n_fresh = int(np.isfinite(bx).sum())
    if n_fresh < MIN_FRESH:
        return None
    raw = nan_xcorr(rx.ravel(), ry.ravel(), MAX_LAG)
    fresh_xc = nan_xcorr(bx, by, MAX_LAG)
    boot_p1 = day_block_boot_xcorr(rx, fx, ry, fy, 1)
    boot_m1 = day_block_boot_xcorr(rx, fx, ry, fy, -1)
    return {
        "n_fresh_pairs": n_fresh,
        "fresh_fraction": round(n_fresh / rx.size, 4),
        "raw_xcorr": {str(k): round(v, 5) if np.isfinite(v) else None for k, v in raw.items()},
        "fresh_xcorr": {str(k): round(v, 5) if np.isfinite(v) else None
                        for k, v in fresh_xc.items()},
        "artifact_share_+1": round(raw[1] - fresh_xc[1], 5)
        if np.isfinite(raw[1]) and np.isfinite(fresh_xc[1]) else None,
        "p_fresh_+1": bootstrap_pvalue(boot_p1, 0.0),
        "p_fresh_-1": bootstrap_pvalue(boot_m1, 0.0),
    }


def main() -> None:
    run_utc = datetime.now(UTC)
    client = HFMarketDataClient()
    random10 = random_stock_universe(
        client.tickers("stock", timeframe="1min", adjustment=ADJ))[:10]

    series_spec: dict[str, tuple[str, str, str | None]] = (
        {"SPY": ("etf", "SPY", ADJ), "SPX": ("index", "SPX", None)}
        | {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}
        | {f"ES[{a}]": ("futures", "ES", a)
           for a in ("contin_adj_ratio", "contin_adj_absolute", "contin_UNadj")}
    )

    pairs: list[dict] = []
    for window, (s, e) in WINDOWS.items():
        grids: dict[str, dict] = {}
        for name, (asset, ticker, adj) in series_spec.items():
            if window == "2006-2007" and (name == "SPX" or name.startswith("ES[")):
                continue
            bars = client.get_bars(asset, ticker, "1min", adj, s, e)
            g = rth_day_grids(bars)
            if len(g) >= MIN_DAYS:
                grids[name] = g
            print(f"{window} {name}: {len(g)} days")
        if "SPY" not in grids:
            continue
        day_list = sorted(grids["SPY"].keys())
        mats = {n: day_matrix(g, day_list) for n, g in grids.items()}
        rx, fx = mats["SPY"]
        for name in mats:
            if name == "SPY":
                continue
            ry, fy = mats[name]
            # convention: x = the hypothesized leader.
            # ES/SPX pairs: x = ES or SPX, y = SPY. Others: x = SPY, y = ticker.
            if name.startswith("ES[") or name == "SPX":
                m = analyze_pair(ry, fy, rx, fx)
            else:
                m = analyze_pair(rx, fx, ry, fy)
            if m is None:
                continue
            bucket = ("es" if name.startswith("ES[") else "index" if name == "SPX"
                      else "sector" if name in SECTORS
                      else "liquid" if name in LIQUID else "random")
            m |= {"pair": f"{name}->SPY" if bucket in ("es", "index") else f"SPY->{name}",
                  "bucket": bucket, "window": window}
            pairs.append(m)
        print(f"{window}: {len(pairs)} pair-cells so far")

    # FDR over all fresh ±1 tests jointly
    tests = [(i, "p_fresh_+1") for i in range(len(pairs))] + \
            [(i, "p_fresh_-1") for i in range(len(pairs))]
    mask = benjamini_hochberg(np.array([pairs[i][k] for i, k in tests]), alpha=0.05)
    for (i, k), rej in zip(tests, mask, strict=True):
        pairs[i][f"fdr_{k[2:]}"] = bool(rej)

    fdr_survivors = [
        {"pair": p["pair"], "window": p["window"], "bucket": p["bucket"],
         "fresh_+1": p["fresh_xcorr"]["1"], "fresh_-1": p["fresh_xcorr"]["-1"],
         "which": [k for k in ("fresh_+1", "fresh_-1") if p[f"fdr_{k}"]]}
        for p in pairs if p.get("fdr_fresh_+1") or p.get("fdr_fresh_-1")
    ]
    summary = {
        "pair_cells": len(pairs),
        "tests": len(tests),
        "fdr_alpha": 0.05,
        "fdr_survivors": fdr_survivors,
        "median_artifact_share_by_bucket": {
            b: round(float(np.median(
                [p["artifact_share_+1"] for p in pairs
                 if p["bucket"] == b and p["artifact_share_+1"] is not None])), 5)
            for b in ("liquid", "sector", "random")
        },
    }

    results = {
        "experiment": "expD_leadlag_scan",
        "run_utc": run_utc.isoformat(),
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "data_source": "hfmarketdata.io",
        "confidence_level": 0,
        "protocol": {"windows": WINDOWS, "adjustment": ADJ, "max_lag": MAX_LAG,
                     "boot": [BOOT_N, BOOT_SEED], "random10": random10,
                     "min_days": MIN_DAYS, "min_fresh": MIN_FRESH},
        "summary": summary,
        "pairs": pairs,
        "client_stats": vars(client.stats) | {"refreshes": list(client.stats.refreshes)},
        "manifest": collect_manifest(),
    }
    out_dir = REPO_ROOT / "results" / "expD_leadlag_scan" / 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(summary, indent=1))


if __name__ == "__main__":
    main()