# SPDX-License-Identifier: MIT
"""YS-EN-0006 · shared experiment runner.

Every config is trained on origins every 3 hours in the training span (adjacent hourly
windows overlap by 167 of 168 hours; `exp_vtt.py check` compares against every hour), tuned/early-stopped on
midnight origins in the validation span, and scored on midnight origins in the test span
(one 24-hour day-ahead forecast per series per day).  Errors are kept per (series, day)
so configs can be compared pairwise on exactly the same days.
"""
from __future__ import annotations

import datetime as dt
import json
import time

import numpy as np

import fc


def run_config(cfg, series, tr, va, te, seeds=(0,), log=print):
    name, kind = cfg["name"], cfg["kind"]
    test_keys = None
    if kind in ("naive", "profile"):
        preds, truth, keys = [], [], []
        for i, s in enumerate(series):
            prof = fc.profile_fit(s, *tr[i]) if kind == "profile" else None
            for t in fc.origins(s, *te[i], step=24):
                preds.append(fc.profile_predict(prof, s, int(t)) if kind == "profile"
                             else fc.naive(s, int(t), cfg["naive"]))
                truth.append(s["y"][t:t + fc.H])
                keys.append((i, int(t)))
        P, Y = np.asarray(preds), np.asarray(truth)
        return {"name": name, "keys": keys, "pred": P[None], "truth": Y, "info": [{}]}

    chans = cfg.get("chans", ["y"])
    scaler = fc.fit_scaler(series, tr, cfg.get("norm", "global"), chans)
    Xtr, Ftr, Ytr, _, _ = fc.build(series, tr, chans, scaler, step=cfg.get("train_step", 3))
    Xva, Fva, Yva, _, _ = fc.build(series, va, chans, scaler, step=24)
    Xte, Fte, Yte, Ote, keys = fc.build(series, te, chans, scaler, step=24)
    R = cfg.get("hist", fc.L)
    Xtr, Xva, Xte = Xtr[:, :, -R:], Xva[:, :, -R:], Xte[:, :, -R:]
    use_cal = cfg.get("cal", True)
    truth = fc.unscale(Yte, Ote, scaler)
    preds, infos = [], []
    if kind == "ridge":
        best = None
        for a in [0.1, 1, 10, 100, 1000, 1e4]:
            m = fc.ridge_fit(Xtr, Ftr, Ytr, a, use_cal)
            v = float(np.abs(fc.ridge_predict(m, Xva, Fva) - Yva).mean())
            if best is None or v < best[0]:
                best = (v, a, m)
        preds.append(fc.unscale(fc.ridge_predict(best[2], Xte, Fte), Ote, scaler))
        infos.append({"alpha": best[1], "best_val": best[0]})
    else:
        for seed in seeds:
            import torch
            torch.manual_seed(seed)
            model = fc.make_model(kind, Xtr.shape[1], use_cal,
                                  n_layers=cfg.get("layers", 8), channels=cfg.get("width", 32),
                                  head=cfg.get("head", "direct"))
            t0 = time.time()
            model, info = fc.train(model, (Xtr, Ftr, Ytr), (Xva, Fva, Yva), seed=seed,
                                   max_epochs=cfg.get("epochs", 40), loss=cfg.get("loss", "l1"))
            info.update(params=fc.n_params(model), seconds=round(time.time() - t0, 1), seed=seed)
            preds.append(fc.unscale(fc.predict(model, Xte, Fte), Ote, scaler))
            infos.append(info)
            log(f"  {name} seed {seed}: {info['epochs']} ep, val {info['best_val']:.4f}, "
                f"{info['params']} params, {info['seconds']} s")
    return {"name": name, "keys": keys, "pred": np.asarray(preds), "truth": truth, "info": infos}


def day_labels(series, keys):
    """Per test day: date (ISO), weekday 0=Mon, workday, public holiday, previous day workday."""
    lab = {"date": [], "dow": [], "workday": [], "holiday": [], "prev_workday": []}
    for i, t in keys:
        s = series[i]
        d = s["start"] + dt.timedelta(t // 24)
        lab["date"].append(d.isoformat())
        lab["dow"].append(d.weekday())
        lab["workday"].append(int(s["workday"][t]))
        lab["holiday"].append(int(d in s["holidays"]))
        lab["prev_workday"].append(int(s["workday"][t - 24]))
    return {k: np.asarray(v) for k, v in lab.items()}


def save(results, path, series):
    """npz with per-config predictions; all configs must share the same test keys."""
    keys = results[0]["keys"]
    out = {"keys": np.asarray(keys), "truth": results[0]["truth"]}
    out.update({"lab__" + k: v for k, v in day_labels(series, keys).items()})
    meta = {}
    for r in results:
        assert r["keys"] == keys, f"{r['name']}: test days differ"
        out["pred__" + r["name"]] = r["pred"].astype(np.float32)
        meta[r["name"]] = r["info"]
    out["meta"] = np.asarray(json.dumps(meta))
    np.savez_compressed(path, **out)