#!/usr/bin/env python3
"""Phase 2 pilot: GARCH(1,1), EGARCH, GJR, APARCH (normal, t) + HAR-RV on a few instruments over the VALIDATION
period only (2012-01-03 .. 2015-12-31), window 1000, refit every 5 days. Writes results/metrics/pilot/*.
Reports: status counts, fit/forecast timings, QLIKE on the common sample, leakage check (future returns perturbed)."""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from wp46.estimation import JobSpec, run_job  # noqa: E402
from wp46.inference import dm_test  # noqa: E402
from wp46.preprocessing import load_panel  # noqa: E402
from wp46.scoring.evaluate import common_sample_mask, score_forecasts  # noqa: E402

SPECS = [("garch11", "normal", "A"), ("garch11", "t", "A"), ("egarch11", "t", "A"), ("gjr11", "t", "A"), ("aparch11", "t", "A"),
         ("nagarch11", "t", "A"), ("histvar22", "empirical", "A"), ("har_rv", "empirical", "B"), ("realgarch11", "normal", "B")]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--symbols", default="SPY,AAPL,CL,EURUSD")
    ap.add_argument("--start", default="2012-01-03")
    ap.add_argument("--end", default="2015-12-31")
    a = ap.parse_args()
    out = ROOT / "results" / "metrics" / "pilot"
    out.mkdir(parents=True, exist_ok=True)
    report = {}
    for sym in a.symbols.split(","):
        panel = load_panel(sym)
        Fs, Gs = [], []
        for mid, dist, tier in SPECS:
            t0 = time.perf_counter()
            F, G = run_job(panel, JobSpec(sym, mid, dist, tier, oos_start=a.start, oos_end=a.end))
            Fs.append(F); Gs.append(G)
            print(f"{sym:7s} {mid:12s} {dist:9s} origins={len(F)} {F['status'].value_counts().to_dict()} fits={G['status'].value_counts().to_dict()} "
                  f"fit_s={G['time_s'].mean():.3f} fc_ms={1000*F['time_fc_s'].mean():.1f} wall={time.perf_counter()-t0:.1f}s", flush=True)
        F = pd.concat(Fs, ignore_index=True); G = pd.concat(Gs, ignore_index=True)
        L = score_forecasts(F, panel, proxies=("rv5ss_on", "r2"), losses=("qlike", "mse"))
        specs = sorted(F["spec_id"].unique())
        cm = common_sample_mask(L, specs)
        Lc = L[cm]
        tab = Lc.groupby(["proxy", "loss", "h", "spec_id"])["value"].mean().unstack("spec_id")
        bench = [s for s in specs if s.startswith("garch11__normal")][0]
        ratio = tab.div(tab[bench], axis=0)
        # leakage check: perturb returns AND realized measures after the last origin -> every forecast of the whole
        # job (same refit schedule) is unchanged
        p2 = panel.copy(); last = F["origin"].max(); m = p2["session"] > last
        p2.loc[m, "r"] *= 3.0; p2.loc[m, "rv_5m_ss"] *= 9.0
        leak_ok = True
        for mid, dist, tier in [("gjr11", "t", "A"), ("har_rv", "empirical", "B")]:
            F1 = F[F.spec_id.str.startswith(f"{mid}__{dist}")].sort_values("origin")
            F2, _ = run_job(p2, JobSpec(sym, mid, dist, tier, oos_start=a.start, oos_end=a.end))
            leak_ok &= bool(np.allclose(F1["h1"].values, F2.sort_values("origin")["h1"].values, equal_nan=True))
        dm = {s: dm_test(Lc[(Lc.spec_id == s) & (Lc.h == 1) & (Lc.proxy == "rv5ss_on") & (Lc.loss == "qlike")].sort_values("origin")["value"].values,
                         Lc[(Lc.spec_id == bench) & (Lc.h == 1) & (Lc.proxy == "rv5ss_on") & (Lc.loss == "qlike")].sort_values("origin")["value"].values, 1)
              for s in specs if s != bench}
        report[sym] = {"n_origins": int(F.groupby("spec_id").size().max()), "common_sample_n_h1": int(cm[(L.h == 1) & (L.proxy == "rv5ss_on") & (L.loss == "qlike")].sum()),
                       "status": F.groupby(["spec_id", "status"]).size().unstack(fill_value=0).to_dict(), "fit_status": G.groupby(["spec_id", "status"]).size().unstack(fill_value=0).to_dict(),
                       "mean_loss_common": tab.round(4).reset_index().to_dict(orient="records"), "ratio_to_garch11_normal": ratio.round(4).reset_index().to_dict(orient="records"),
                       "dm_vs_garch11_normal_h1_qlike": {k: {kk: (round(vv, 4) if isinstance(vv, float) else vv) for kk, vv in v.items()} for k, v in dm.items()},
                       "leakage_check_passed": leak_ok, "timing": G.groupby("spec_id")["time_s"].agg(["mean", "median", "max"]).round(4).to_dict()}
        L.to_parquet(out / f"losses_{sym}.parquet", index=False)
        F.to_parquet(out / f"forecasts_{sym}.parquet", index=False)
        print(sym, "QLIKE ratio to GARCH(1,1)-N, proxy rv5ss_on:"); print(ratio.loc[("rv5ss_on", "qlike")].round(3).to_string())
        print("leakage check:", leak_ok)
    (out / "pilot_report.json").write_text(json.dumps(report, indent=1, default=str))


if __name__ == "__main__":
    main()