"""Glue: grids for a symbol set, chronological session splits, origin indices, baseline fitting and scoring tables."""
from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np
import pandas as pd

from .data.windows import SymbolGrid, WindowConfig, OriginIndex, make_example, trailing_sigma
from .eval import scores as S
from .eval.calibration import pit_histogram, reliability
from .eval.statistics import diebold_mariano, cluster_bootstrap_mean


# ------------------------------------------------------------------------------------------------ grids & splits
def grid_dirs(cache: Path, start: str, end: str, symbols: list[str]) -> list[Path]:
    return [d for d in (cache / "grids" / f"{start}_{end}" / s for s in symbols) if d.exists()]


def load_grids(cache: Path, start: str, end: str, symbols: list[str], mmap: bool = True) -> list[SymbolGrid]:
    return [SymbolGrid.load(d, mmap=mmap) for d in grid_dirs(cache, start, end, symbols)]


def session_split(g: SymbolGrid, bounds: dict[str, str], warmup: int = 21) -> dict[str, np.ndarray]:
    """Session index sets per split from date bounds {train_end, val_end, test_end[, final_end]}. The first `warmup`
    sessions are never origins; each split uses sessions strictly after the previous bound."""
    dates = np.array(g.meta["session_dates"], dtype="datetime64[D]")
    idx = np.arange(len(dates))
    prev = np.datetime64("1900-01-01")
    out = {}
    for name in ("train", "val", "test", "final"):
        key = f"{name}_end"
        if key not in bounds:
            continue
        end = np.datetime64(bounds[key])
        sel = idx[(dates > prev) & (dates <= end) & (idx >= warmup)]
        out[name] = sel
        prev = end
    return out


def build_indices(dirs: list[Path], bounds: dict[str, str], wcfg: WindowConfig, warmup: int, train_per_session: int,
                  eval_stride: int, seed: int, embargo_sessions: int = 1) -> dict[str, OriginIndex]:
    """Train origins: random `train_per_session` per session; eval origins: fixed stride. An embargo of
    `embargo_sessions` sessions is removed at the start of val/test/final so that no training target overlaps."""
    grids = [SymbolGrid.load(d) for d in dirs]
    splits = [session_split(g, bounds, warmup) for g in grids]
    out = {}
    for name in splits[0].keys():
        sets = []
        for sp in splits:
            s = sp[name]
            if name != "train" and len(s) > embargo_sessions:
                s = s[embargo_sessions:]
            sets.append(s)
        if name == "train":
            out[name] = OriginIndex(dirs, wcfg, sets, stride=1, max_per_session=train_per_session, seed=seed, grids=grids)
        else:
            out[name] = OriginIndex(dirs, wcfg, sets, stride=eval_stride, grids=grids)
    return out


# ------------------------------------------------------------------------------------------------ baselines
def origin_table(index: OriginIndex, horizons: list[int], garch_params: dict | None = None, workers: int = 8) -> pd.DataFrame:
    """Realised targets, sigma, metadata and baseline features for every origin (single causal pass, parallel)."""
    from .baselines.features import origin_frame
    return origin_frame(index, horizons, garch_params, workers)


def samples_from_params(p: dict, n: int, seed: int = 0) -> np.ndarray:
    """Student-t / Gaussian samples [n, N] from arrays mu, scale, nu."""
    rng = np.random.default_rng(seed)
    mu, sc, nu = (np.asarray(p[k], dtype=np.float64) for k in ("mu", "scale", "nu"))
    N = len(mu)
    z = rng.standard_normal((n, N))
    nu_b = np.broadcast_to(nu, (N,))
    fin = np.isfinite(nu_b)
    g = np.ones((n, N))
    if fin.any():
        g[:, fin] = rng.chisquare(nu_b[fin], (n, fin.sum())) / nu_b[fin]
    return mu[None] + sc[None] * z / np.sqrt(g)


# ------------------------------------------------------------------------------------------------ scoring
LEVELS = np.array([0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975])


def score_horizon(y: np.ndarray, samples: np.ndarray, rv: np.ndarray, rv_hat: np.ndarray, sigma: np.ndarray,
                  mu_hat: np.ndarray | None = None, tail_k: float = 2.5, H: int = 60,
                  params: dict | None = None) -> dict[str, np.ndarray]:
    """Per-origin scores at one horizon from cumulative-return samples [n, N]."""
    out = {"crps": S.crps_samples(y, samples), "wis": S.weighted_interval_score(y, samples),
           "qloss": S.quantile_loss(y, S.quantiles_from_samples(samples, LEVELS), LEVELS),
           "mae": np.abs(y - np.median(samples, 0)), "mse": (y - samples.mean(0)) ** 2,
           "qlike": S.qlike(rv, rv_hat), "pit": S.pit_samples(y, samples), "sharpness": np.quantile(samples, 0.95, 0) - np.quantile(samples, 0.05, 0)}
    thr = tail_k * sigma * np.sqrt(H)
    event = np.abs(y) > thr
    p_tail = (np.abs(samples) > thr[None]).mean(0)
    out["brier_tail"] = S.brier(event, p_tail); out["tail_event"] = event.astype(float)
    mu = samples.mean(0) if mu_hat is None else mu_hat
    out["dir_hit"] = ((np.sign(mu) == np.sign(y)) & (y != 0)).astype(float)
    out["crps_scaled"] = out["crps"] / (sigma * np.sqrt(H))
    if params is not None and np.all(np.isfinite(params["nu"])):
        out["logscore"] = S.log_score_student_t(y, params["mu"], params["scale"], params["nu"])
        out["crps_closed"] = S.crps_student_t(y, params["mu"], params["scale"], params["nu"])
    return out


def summarise(per: dict[str, np.ndarray], meta: pd.DataFrame, H: int) -> dict:
    """Aggregate per-origin scores with a cluster bootstrap by trading day."""
    cl = pd.factorize(meta["session_date"])[0]
    out = {"H": H, "n": int(len(meta)), "n_days": int(len(np.unique(cl)))}
    for k in ("crps", "crps_scaled", "wis", "qloss", "mae", "mse", "qlike", "brier_tail", "sharpness", "logscore", "crps_closed"):
        if k in per:
            cb = cluster_bootstrap_mean(per[k], cl, n_boot=500)
            out[k] = cb["mean"]; out[f"{k}_lo"] = cb["lo"]; out[f"{k}_hi"] = cb["hi"]
    out["dir_acc"] = float(np.nanmean(per["dir_hit"][meta[f"cum_{H}"].values != 0]))
    out["tail_rate"] = float(per["tail_event"].mean())
    for lv in (0.5, 0.8, 0.95):
        a = 1 - lv
        pit = per["pit"]
        out[f"cov_{int(lv * 100)}"] = float(((pit >= a / 2) & (pit <= 1 - a / 2)).mean())
    ph = pit_histogram(per["pit"])
    out["pit_ks"] = ph["ks_stat"]; out["pit_chi2_p"] = ph["chi2_p"]; out["pit_var"] = ph["var"]
    return out


def compare(per_a: np.ndarray, per_b: np.ndarray, meta: pd.DataFrame, H: int, eval_stride: int) -> dict:
    """Model a vs b on a loss: DM with HAC lag = ceil(H/stride) overlap and day-cluster bootstrap CI of the mean diff."""
    order = pd.to_datetime(meta["session_date"]).values.astype("datetime64[m]").astype(np.int64) * 1000 + meta["m"].values
    # chronological order across symbols: sort by (date, slot); symbols interleave within a slot -> cluster by day
    lag = int(np.ceil(H / max(1, eval_stride))) * max(1, meta["symbol"].nunique())
    dm = diebold_mariano(per_a, per_b, order, lag=min(lag, len(per_a) // 4))
    cl = pd.factorize(meta["session_date"])[0]
    cb = cluster_bootstrap_mean(per_a - per_b, cl, n_boot=1000)
    return {"H": H, "diff": cb["mean"], "lo": cb["lo"], "hi": cb["hi"], "dm_t": dm["t"], "dm_p": dm["p"], "lag": dm["lag"],
            "n": int(len(per_a)), "rel_diff": float(cb["mean"] / max(1e-12, per_b.mean()))}