"""Run one experiment config end to end: indices -> baselines -> neural variants -> predictions -> scores -> tables.
LightGBM (baselines) and torch (neural) cannot share a process on macOS (duplicate libomp: segfault / Cholesky
deadlock, verified 2026-11-02), so the stages run as separate subprocesses and communicate through files.
Usage: run_experiment.py --config configs/pilot.yaml [--stage all|baselines|neural|tables] [--symbols-from pilot|feasibility|all]
       [--eval-split test|final] [--only-models m1,m2] [--skip-baselines] [--skip-train] [--exp-name NAME] [--heldout] [--models-from DIR]"""
from __future__ import annotations

import sys

if "baselines" in sys.argv:          # LightGBM must be the first native library loaded in the baselines process
    import lightgbm  # noqa: F401

import argparse
import json
import subprocess
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 chronos_x.config import load_config  # noqa: E402
from chronos_x import pipeline as P  # noqa: E402
from chronos_x.data.windows import WindowConfig  # noqa: E402
from chronos_x.inference.serialization import write_json  # noqa: E402


def parse():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/pilot.yaml"); ap.add_argument("--stage", default="all")
    ap.add_argument("--symbols-from", default=None); ap.add_argument("--eval-split", default="test")
    ap.add_argument("--only-models", default=None); ap.add_argument("--skip-baselines", action="store_true")
    ap.add_argument("--skip-train", action="store_true", help="reuse saved checkpoints"); ap.add_argument("--exp-name", default=None)
    ap.add_argument("--heldout", action="store_true", help="evaluate on held-out instruments (zero-shot) with saved models")
    ap.add_argument("--models-from", default=None, help="experiment directory holding trained models to reuse")
    return ap.parse_args()


def setup(a):
    from chronos_x import experiment as E
    cfg = load_config(a.config)
    exp = a.exp_name or cfg["name"]
    out = ROOT / "results" / "metrics" / exp; out.mkdir(parents=True, exist_ok=True)
    man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())
    which = a.symbols_from or cfg["universe"].get("symbols_from", "pilot")
    roles = ("heldout_instrument",) if a.heldout else ("pooled",)
    rows = E.symbols_for(man, which, roles)
    d = cfg["data"]
    dirs = P.grid_dirs(Path(cfg["paths"]["cache"]), d["start"], d["end"], [r["symbol"] for r in rows])
    wcfg = WindowConfig(**cfg["windows"]); ev = cfg["eval"]; tr = cfg["train"]
    idx = P.build_indices(dirs, cfg["splits"], wcfg, d["warmup_sessions"], tr["origins_per_session"], ev["stride"], tr["seed"],
                          embargo_sessions=cfg["splits"].get("embargo_sessions", 1))
    tag = a.eval_split + ("_heldout" if a.heldout else "")
    empty = [k for k, v in idx.items() if len(v) == 0 and k in ("train", a.eval_split)]
    if empty:
        raise SystemExit(f"no eligible origins for split(s) {empty}: check splits/data window in {a.config}")
    return cfg, exp, out, man, rows, dirs, idx, tag


def stage_baselines(a):
    from chronos_x import experiment as E
    cfg, exp, out, man, rows, dirs, idx, tag = setup(a)
    ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]
    idx_eval = idx[a.eval_split]
    sizes = {k: len(v) for k, v in idx.items()}
    print(f"[{exp}:baselines] {len(dirs)} grids, origins {sizes}", flush=True)
    write_json({"config": cfg, "symbols": [r["symbol"] for r in rows], "origins": sizes, "eval_split": a.eval_split, "heldout": a.heldout},
               out / f"run_meta_{tag}.json")
    t0 = time.time()
    which_b = tuple(cfg.get("baselines", ["naive", "har", "garch", "trees"]))
    garch_params = E.fit_garch_params(idx["train"]) if "garch" in which_b else None
    tab = P.origin_table(idx_eval, horizons, garch_params, workers=tr["num_workers"])
    tab.to_csv(out / f"origins_{tag}.csv", index=False)
    print(f"[{exp}:baselines] eval origin frame {len(tab)} rows in {time.time() - t0:.0f}s", flush=True)
    t0 = time.time()
    tab_tr = P.origin_table(idx["train"], horizons, None, workers=tr["num_workers"])
    tab_tr.to_csv(out / "origins_train.csv", index=False)
    preds = E.fit_and_predict_baselines(tab_tr, tab, horizons, out, tr["seed"], which_b)
    if garch_params is not None:
        write_json({str(i): p for i, p in garch_params.items()}, out / "garch_params.json")
    for name, pH in preds.items():
        E.score_predictions(f"{name}_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], param_preds=pH, universe=man)
    print(f"[{exp}:baselines] done in {time.time() - t0:.0f}s: {list(preds)}", flush=True)


def stage_neural(a):
    import resource
    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
    resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard))       # memory-mapped grids x workers
    import torch
    torch.multiprocessing.set_sharing_strategy("file_system")
    from chronos_x import experiment as E
    from chronos_x.inference.predict import predict_index
    from chronos_x.inference.serialization import save_predictions
    from chronos_x.train import checkpoint as ck
    from chronos_x.eval.statistics import cluster_bootstrap_mean
    cfg, exp, out, man, rows, dirs, idx, tag = setup(a)
    ev, tr, horizons = cfg["eval"], cfg["train"], cfg["eval"]["horizons"]
    idx_eval = idx[a.eval_split]; n_grids = len(dirs)
    f_tab = out / f"origins_{tag}.csv"
    if f_tab.exists():
        tab = pd.read_csv(f_tab)
    else:
        tab = P.origin_table(idx_eval, horizons, None, workers=tr["num_workers"]); tab.to_csv(f_tab, index=False)
    print(f"[{exp}:neural] {n_grids} grids, origins { {k: len(v) for k, v in idx.items()} }, device {E.pick_device()}", flush=True)
    only = set(a.only_models.split(",")) if a.only_models else None
    for spec in cfg.get("models", []):
        name = spec["name"]
        if only and name not in only:
            continue
        mdir = (Path(a.models_from) / name) if a.models_from else out / name
        mdir.mkdir(parents=True, exist_ok=True)
        t0 = time.time()
        if a.skip_train or a.heldout or a.models_from:
            if spec.get("kwargs", {}).get("ticker_embedding") and a.heldout:
                print(f"[{exp}:neural] {name}: ticker embedding cannot be used zero-shot, skipping", flush=True); continue
            model = E.build_model(spec, 0 if a.heldout else n_grids)
            ck.load(mdir / "best.pt", model)
            if getattr(model, "copula", None) is not None and (mdir / "copula" / "best.pt").exists():
                ck.load(mdir / "copula" / "best.pt", model)
            summ = json.loads((mdir / "train_summary.json").read_text()) if (mdir / "train_summary.json").exists() else {}
        else:
            model, summ = E.train_variant(spec, idx["train"], idx["val"], cfg, mdir, n_grids)
            ck.load(mdir / "best.pt", model)
            if getattr(model, "copula", None) is not None:
                ck.load(mdir / "copula" / "best.pt", model)
            write_json(summ, mdir / "train_summary.json")
        joint = getattr(model, "copula", None) is not None
        kw = dict(n_samples=ev["n_samples"], batch_size=tr.get("eval_batch_size", 128), feature_set=spec.get("feature_set", "full"),
                  num_workers=tr["num_workers"], seed=tr["seed"], keep_paths=spec.get("keep_paths", False))
        res = predict_index(model, idx_eval, horizons, shuffle_minutes=spec.get("shuffle_minutes_eval", False), joint=joint, **kw)
        mtag = f"{name}_{tag}"
        save_predictions(res, out / f"pred_{mtag}.npz")
        E.score_predictions(mtag, tab, horizons, out, ev["n_samples"], tr["seed"], neural=res, universe=man)
        if joint:
            res_i = predict_index(model, idx_eval, horizons, joint=False, **kw)
            E.score_predictions(f"{name}_indep_{tag}", tab, horizons, out, ev["n_samples"], tr["seed"], neural=res_i, universe=man)
            if spec.get("keep_paths"):
                cl = pd.factorize(tab["session_date"])[0]; ps = {}
                for lab, r in (("joint", res), ("indep", res_i)):
                    ps[lab] = {k: cluster_bootstrap_mean(r[k], cl, n_boot=500) for k in ("energy", "variogram", "mdd_crps", "rv_crps")}
                ps["joint_minus_indep"] = {k: cluster_bootstrap_mean(res[k] - res_i[k], cl, n_boot=1000) for k in ("energy", "variogram", "mdd_crps", "rv_crps")}
                write_json(ps, out / f"path_scores_{mtag}.json")
        print(f"[{exp}:neural] {name}: {time.time() - t0:.0f}s, params={model.n_params()}, train={summ.get('steps')} steps, best_val={summ.get('best_val')}", flush=True)


def stage_tables(a):
    from chronos_x import experiment as E
    cfg = load_config(a.config); exp = a.exp_name or cfg["name"]
    out = ROOT / "results" / "metrics" / exp
    tag = a.eval_split + ("_heldout" if a.heldout else "")
    horizons = cfg["eval"]["horizons"]; ev = cfg["eval"]
    tab = pd.read_csv(out / f"origins_{tag}.csv")
    scored = {}
    for f in sorted(out.glob(f"summary_*_{tag}.json")):
        name = f.name[len("summary_"):-len(f"_{tag}.json")]
        per = np.load(out / f"per_origin_{name}_{tag}.npz")
        scored[name] = {"summary": {int(H): v for H, v in json.loads(f.read_text()).items()},
                        "per": {H: {k[:-len(f"_{H}")]: per[k] for k in per.files if k.endswith(f"_{H}")} for H in horizons}}
    lb = E.leaderboard(scored, horizons); lb.to_csv(out / f"leaderboard_{tag}.csv", index=False)
    ref = cfg.get("reference", "naive_rw")
    if ref in scored:
        for metric in ("crps", "wis", "qlike"):
            E.comparisons(scored, tab, horizons, ev["stride"], metric, ref).to_csv(out / f"compare_{metric}_vs_{ref}_{tag}.csv", index=False)
    cols = ["model", "H", "crps_scaled", "wis", "qlike", "cov_50", "cov_80", "cov_95", "pit_ks", "dir_acc"]
    print(lb[[c for c in cols if c in lb.columns]].to_string(index=False, float_format=lambda x: f"{x:.5g}"))


def main() -> None:
    a = parse()
    if a.stage == "baselines":
        stage_baselines(a)
    elif a.stage == "neural":
        stage_neural(a)
    elif a.stage == "tables":
        stage_tables(a)
    else:
        passthrough = [x for x in sys.argv[1:] if x not in ("--stage", "all")]
        if "--stage" in sys.argv:
            i = sys.argv.index("--stage"); passthrough = [x for j, x in enumerate(sys.argv[1:], 1) if j not in (i, i + 1)]
        if not a.skip_baselines:
            subprocess.run([sys.executable, __file__, "--stage", "baselines", *passthrough], check=True)
        subprocess.run([sys.executable, __file__, "--stage", "neural", *passthrough], check=True)
        stage_tables(a)


if __name__ == "__main__":
    main()