"""End-to-end experiment: grids -> splits -> baselines + neural variants -> predictions -> scores -> tables."""
from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from .config import load_config
from .data.windows import WindowConfig, OriginIndex, make_example
from .models.model import MRPT, MRPTConfig
from .baselines.har import HARRV
from .baselines.trees import RidgeBaseline, ElasticNetBaseline, LightGBMBaseline
from .baselines.garch import five_minute_returns
from .baselines.neural import MLPBaseline, TCNBaseline
from .train.trainer import train, pick_device
from .train.sampling import OriginDataset
from .inference.predict import predict_index
from .inference.serialization import save_predictions, write_json
from . import pipeline as P

ROOT = Path(__file__).resolve().parents[2]


def symbols_for(man: dict, which: str, roles=("pooled",)) -> list[dict]:
    inst = {r["symbol"]: r for r in man["instruments"]}
    if which == "pilot":
        syms = man["pilot_symbols"]
    elif which == "feasibility":
        syms = man["feasibility_symbols"]
    else:
        syms = [s for s, r in inst.items() if r["role"] in roles]
    return [inst[s] for s in syms if inst[s]["role"] in roles]


# ------------------------------------------------------------------------------------------------ baselines
def _garch_fit(r5: np.ndarray) -> dict | None:
    from arch import arch_model
    import warnings
    if len(r5) < 500:
        return None
    y = r5 * 1e4
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        res = arch_model(y, mean="Zero", vol="GARCH", p=1, q=1, dist="t").fit(disp="off", show_warning=False)
    p = res.params
    return {"omega": float(p["omega"]) / 1e8, "alpha": float(p["alpha[1]"]), "beta": float(p["beta[1]"]),
            "nu": float(np.clip(p["nu"], 2.5, 50)), "n_obs": int(len(r5))}


def fit_garch_params(idx_train: OriginIndex) -> dict[int, dict]:
    """GARCH(1,1)-t per instrument on the 5-minute returns of its training sessions."""
    out = {}
    for i in np.unique(idx_train.rows[:, 0]):
        s_tr = idx_train.rows[idx_train.rows[:, 0] == i, 1]
        pr = _garch_fit(five_minute_returns(idx_train.grid(int(i)), int(s_tr.min()), int(s_tr.max())))
        if pr is not None:
            out[int(i)] = pr
    return out


def fit_and_predict_baselines(tab_tr: pd.DataFrame, tab_ev: pd.DataFrame, horizons: list[int], out_dir: Path, seed: int,
                              which: tuple[str, ...]) -> dict[str, dict[int, dict]]:
    """Baselines from the origin frames (training frame for fitting, evaluation frame for prediction)."""
    from .baselines.features import TREE_FEATURES, HAR_FEATURES
    t0 = time.time()
    N = len(tab_ev); sig = tab_ev["sigma"].values
    preds: dict[str, dict[int, dict]] = {}
    inf = np.full(N, np.inf); zero = np.zeros(N)
    if "naive" in which:
        preds["naive_rw"] = {H: {"mu": zero, "scale": sig * np.sqrt(H), "nu": inf, "var": sig ** 2 * H} for H in horizons}
        preds["seasonal_vol"] = {H: {"mu": zero, "scale": np.sqrt(tab_ev[f"seas_var_{H}"].values), "nu": inf,
                                     "var": tab_ev[f"seas_var_{H}"].values} for H in horizons}
        nu = 5.0
        preds["ewma"] = {H: {"mu": zero, "scale": np.sqrt(tab_ev["ewma_var1"].values * H * (nu - 2) / nu), "nu": np.full(N, nu),
                             "var": tab_ev["ewma_var1"].values * H} for H in horizons}
    if "har" in which:
        har = HARRV(); preds["har_rv"] = {}
        Xtr = tab_tr[HAR_FEATURES].values; Xev = tab_ev[HAR_FEATURES].values
        for H in horizons:
            har.fit(Xtr, tab_tr["m"].values, tab_tr[f"rv_{H}"].values, tab_tr[f"cum_{H}"].values, H)
            preds["har_rv"][H] = har.predict_batch(Xev, tab_ev["m"].values, H)
    if "garch" in which and f"garch_var_{horizons[0]}" in tab_ev.columns:
        preds["garch_t"] = {}
        for H in horizons:
            var = tab_ev[f"garch_var_{H}"].values.astype(float); nu = tab_ev["garch_nu"].values.astype(float)
            miss = ~np.isfinite(var) | (var <= 0)
            var = np.where(miss, sig ** 2 * H, var); nu = np.where(np.isfinite(nu), nu, 6.0)
            preds["garch_t"][H] = {"mu": zero, "scale": np.sqrt(var * (nu - 2) / nu), "nu": nu, "var": var}
    if "trees" in which:
        Xtr = tab_tr[TREE_FEATURES].values; Xev = tab_ev[TREE_FEATURES].values
        for name, model in (("ridge", RidgeBaseline()), ("elastic_net", ElasticNetBaseline()), ("lightgbm", LightGBMBaseline(seed=seed))):
            preds[name] = {}
            for H in horizons:
                model.fit(Xtr, tab_tr[f"cum_{H}"].values, tab_tr[f"rv_{H}"].values, tab_tr["sigma"].values, H)
                preds[name][H] = model.predict_batch(Xev, sig, H)
    write_json({"wall_s": time.time() - t0, "n_train": int(len(tab_tr)), "n_eval": N, "baselines": list(preds)}, out_dir / "baselines_meta.json")
    return preds


# ------------------------------------------------------------------------------------------------ neural
def build_model(spec: dict, n_symbols: int) -> torch.nn.Module:
    kind = spec.get("kind", "mrpt")
    if kind == "mlp":
        return MLPBaseline()
    if kind == "tcn":
        return TCNBaseline()
    kw = dict(spec.get("kwargs", {}))
    if kw.pop("ticker_embedding", False):
        kw["n_symbols"] = n_symbols
    return MRPT(MRPTConfig.preset(spec.get("preset", "small"), **kw))


def train_variant(spec: dict, idx_train: OriginIndex, idx_val: OriginIndex, cfg: dict, out_dir: Path, n_symbols: int) -> tuple[torch.nn.Module, dict]:
    tcfg = cfg["train"]
    model = build_model(spec, n_symbols)
    ds_tr = OriginDataset(idx_train, spec.get("feature_set", "full"), spec.get("shuffle_minutes", False))
    ds_va = OriginDataset(idx_val, spec.get("feature_set", "full"))
    summ = train(model, ds_tr, ds_va, out_dir, batch_size=tcfg["batch_size"], max_steps=spec.get("max_steps", tcfg["max_steps"]),
                 lr=tcfg["lr"], warmup_steps=tcfg["warmup_steps"], weight_decay=tcfg["weight_decay"], grad_clip=tcfg["grad_clip"],
                 eval_every=tcfg["eval_every"], patience=tcfg.get("patience", 5), seed=spec.get("seed", tcfg["seed"]),
                 num_workers=tcfg["num_workers"], aux_vol_weight=tcfg.get("aux_vol_weight", 0.0),
                 val_max_batches=tcfg.get("val_max_batches", 50))
    # Stage B: copula head trained after the marginals, marginals frozen
    if getattr(model, "copula", None) is not None:
        from .train import checkpoint as ck
        ck.load(out_dir / "best.pt", model)
        summ_b = train(model, ds_tr, ds_va, out_dir / "copula", batch_size=tcfg["batch_size"], max_steps=spec.get("copula_steps", 300),
                       lr=tcfg["lr"], warmup_steps=20, weight_decay=0.0, grad_clip=tcfg["grad_clip"], eval_every=100, patience=3,
                       seed=spec.get("seed", tcfg["seed"]), num_workers=tcfg["num_workers"], copula_only=True, resume=False)
        summ["copula"] = summ_b
    return model, summ


# ------------------------------------------------------------------------------------------------ scoring
def score_predictions(name: str, tab: pd.DataFrame, horizons: list[int], out_dir: Path, n_samples: int, seed: int,
                      neural: dict | None = None, param_preds: dict | None = None, universe: dict | None = None) -> dict:
    """Per-origin scores at each horizon (saved) and summaries. Either `neural` (predict_index output) or
    `param_preds` {H: {mu, scale, nu, var}} must be given."""
    from .eval.subgroups import add_subgroups
    meta = add_subgroups(tab, universe) if universe else tab
    summaries, per_all = {}, {}
    for H in horizons:
        y = tab[f"cum_{H}"].values; rv = tab[f"rv_{H}"].values; sigma = tab["sigma"].values
        if neural is not None:
            samples = neural[f"cum_samples_{H}"]; rv_hat = neural[f"rv_samples_{H}"].mean(0)
            # consistency: realised from the predictor must equal the origin table
            assert np.allclose(neural[f"cum_{H}"], y, atol=1e-6), "origin/prediction misalignment"
            mu_hat = samples.mean(0); params = None
        else:
            p = param_preds[H]; samples = P.samples_from_params(p, n_samples, seed); rv_hat = p["var"]; mu_hat = p["mu"]
            params = p if np.all(np.isfinite(p["nu"])) else None
        per = P.score_horizon(y, samples, rv, rv_hat, sigma, mu_hat, H=H, params=params)
        per_all[H] = per
        summaries[H] = P.summarise(per, meta, H)
    np.savez_compressed(out_dir / f"per_origin_{name}.npz", **{f"{k}_{H}": v for H, per in per_all.items() for k, v in per.items()})
    write_json(summaries, out_dir / f"summary_{name}.json")
    return {"summary": summaries, "per": per_all}


def comparisons(scored: dict[str, dict], tab: pd.DataFrame, horizons: list[int], stride: int, metric: str, reference: str) -> pd.DataFrame:
    rows = []
    for name, sc in scored.items():
        if name == reference:
            continue
        for H in horizons:
            r = P.compare(sc["per"][H][metric], scored[reference]["per"][H][metric], tab, H, stride)
            rows.append({"model": name, "reference": reference, "metric": metric, **r})
    return pd.DataFrame(rows)


def leaderboard(scored: dict[str, dict], horizons: list[int]) -> pd.DataFrame:
    rows = []
    for name, sc in scored.items():
        for H in horizons:
            s = sc["summary"][H]
            rows.append({"model": name, "H": H, **{k: s[k] for k in s if k not in ("H",)}})
    return pd.DataFrame(rows)