"""Batched prediction over an OriginIndex: Student-t parameters per future minute, Monte Carlo cumulative-return
samples at each horizon (Stage A: independent minutes; Stage B: copula-correlated), plus metadata for scoring."""
from __future__ import annotations

import numpy as np
import torch
from torch.utils.data import DataLoader

from ..data.windows import OriginIndex
from ..train.sampling import OriginDataset, collate
from ..train.trainer import pick_device, to_device
from ..models.student_t import sample_student_t
from ..models.joint_path import samples_from_copula


@torch.no_grad()
def predict_index(model, index: OriginIndex, horizons: list[int], n_samples: int = 256, batch_size: int = 128,
                  feature_set: str = "full", shuffle_minutes: bool = False, device: str | None = None, num_workers: int = 4,
                  joint: bool = False, seed: int = 0, keep_paths: bool = False) -> dict[str, np.ndarray]:
    dev = pick_device(device)
    model.to(dev).eval()
    ds = OriginDataset(index, feature_set, shuffle_minutes)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate)
    gen = torch.Generator(device=dev); gen.manual_seed(seed)
    out = {k: [] for k in ("mu", "scale", "nu", "y", "y_mask", "sigma", "origin", "sym")}
    cum = {H: [] for H in horizons}; realised = {H: [] for H in horizons}; rv = {H: [] for H in horizons}
    rv_pred = {H: [] for H in horizons}
    path_sc = {k: [] for k in ("energy", "variogram", "mdd_crps", "rv_crps")}
    for b in loader:
        bd = to_device(b, dev)
        o = model(bd)
        if "quantiles" in o:
            raise NotImplementedError("sample-based scoring needs a parametric head")
        mu, sc, nu = o["mu"], o["scale"], o["nu"]
        if joint and getattr(model, "copula", None) is not None:
            z = model.copula.sample_z(o["hidden"], n_samples, gen)
            S = samples_from_copula(z, mu, sc, nu)                     # [n, B, H] cpu
        else:
            S = sample_student_t({"mu": mu, "scale": sc, "nu": nu}, n_samples, gen).cpu()
        S = S.numpy()
        ym = b["y_mask"].numpy().astype(bool); y = b["y"].numpy()
        for H in horizons:
            Sm = S[:, :, :H] * ym[None, :, :H]       # missing future minutes contribute nothing (same as realised)
            cum[H].append(Sm.sum(-1)); realised[H].append((y[:, :H] * ym[:, :H]).sum(-1))
            rv[H].append(((y[:, :H] * ym[:, :H]) ** 2).sum(-1)); rv_pred[H].append((Sm ** 2).sum(-1))
        if keep_paths:
            from ..eval.scores import energy_score, variogram_score, crps_samples, path_stats
            Y = y * ym; Sp = S * ym[None]
            path_sc["energy"].append(energy_score(Y, Sp)); path_sc["variogram"].append(variogram_score(Y, Sp))
            st = path_stats(Sp); ob = path_stats(Y[None])
            path_sc["mdd_crps"].append(crps_samples(ob["mdd"][0], st["mdd"])); path_sc["rv_crps"].append(crps_samples(ob["rv"][0], st["rv"]))
        for k in ("mu", "scale", "nu"):
            out[k].append(o[k].float().cpu().numpy())
        out["y"].append(y); out["y_mask"].append(ym); out["sigma"].append(b["sigma"].numpy())
        out["origin"].append(b["origin"].numpy()); out["sym"].append(b["sym"].numpy())
    res = {k: np.concatenate(v) for k, v in out.items()}
    for H in horizons:
        res[f"cum_samples_{H}"] = np.concatenate(cum[H], axis=1)      # [n, N]
        res[f"cum_{H}"] = np.concatenate(realised[H]); res[f"rv_{H}"] = np.concatenate(rv[H])
        res[f"rv_samples_{H}"] = np.concatenate(rv_pred[H], axis=1)
    if keep_paths:
        for k, v in path_sc.items():
            res[k] = np.concatenate(v)
    return res