#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task B robustness (protocol §12): rolling 24-origin training windows vs expanding, three seeds for the neural models,
exclusion of thin cell-months (n_last < 10), DSGFN bandwidths 30 / 120 km, winsorised targets. Torch-family process.
Writes results/robustness_b/predictions_robustness.parquet and summary.csv."""
from __future__ import annotations

import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import numpy as np
import polars as pl

from wp23 import config as C
from wp23.data import month_add
from wp23.models_b import AR2, Momentum, NoChange, RidgePanel, TAUS
from wp23.models_b_torch import DSGFN, QANSS, TFTLite

OUT = C.RESULTS / "robustness_b"; OUT.mkdir(parents=True, exist_ok=True)


def run(panel, models, tag, h=3, window=None, min_n=None, winsor=False, refit=3):
    rows = []
    ph = panel.filter(pl.col("h") == h)
    if min_n: ph = ph.filter(pl.col("n_last") >= min_n)
    if winsor:
        q1, q99 = ph["target"].quantile(0.01), ph["target"].quantile(0.99); ph = ph.with_columns(pl.col("target").clip(q1, q99))
    origins = sorted(ph["origin"].unique().to_list())
    for k in range(0, len(origins), refit):
        o = origins[k]; test_o = origins[k:k + refit]
        tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)
        if window: tr = tr.filter(pl.col("origin") > month_add(o, -window))
        if tr["origin"].n_unique() < 10: continue
        te = ph.filter(pl.col("origin").is_in(test_o))
        for m in models:
            t0 = time.time()
            try: m.fit(tr); p = m.predict(te); Q = m.predict_quantiles(te)
            except Exception as ex: print("FAILED", tag, o, m.name, repr(ex), flush=True); continue
            rows.append(te.select(["class", "level", "geo_id", "origin", "h", "target", "n_last"]).with_columns(pl.lit(tag).alias("variant"), pl.lit(m.name).alias("model"), pl.Series("pred", np.asarray(p, float)), pl.Series("q05", np.asarray(Q[:, 0], float)), pl.Series("q95", np.asarray(Q[:, 4], float))))
            print(tag, o, m.name, f"{time.time() - t0:.1f}s", flush=True)
    return pl.concat(rows) if rows else None


def main() -> None:
    panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet"); parts = []
    base = lambda: [NoChange(), Momentum(), AR2(), RidgePanel(), QANSS(), TFTLite(), DSGFN()]
    parts.append(run(panel, base(), "expanding"))
    parts.append(run(panel, base(), "rolling_24", window=24))
    parts.append(run(panel, base(), "min_n_10", min_n=10))
    parts.append(run(panel, base(), "winsorised", winsor=True))
    for s in (1, 2):
        parts.append(run(panel, [QANSS(seed=s), TFTLite(seed=s), DSGFN(seed=s)], f"seed_{s}"))
    parts.append(run(panel, [DSGFN(bandwidth_km=30.0), DSGFN(bandwidth_km=120.0)], "bandwidth"))
    # label the bandwidth variants
    out = []
    for p in parts:
        if p is not None: out.append(p)
    allp = pl.concat(out); allp.write_parquet(OUT / "predictions_robustness.parquet")
    s = allp.with_columns((pl.col("pred") - pl.col("target")).abs().alias("ae"), ((pl.col("target") >= pl.col("q05")) & (pl.col("target") <= pl.col("q95"))).cast(pl.Float64).alias("cov90")).group_by(["variant", "model", "level"]).agg(pl.len().alias("n"), pl.col("ae").mean().alias("mae"), pl.col("cov90").mean().alias("cov90")).sort(["level", "variant", "mae"])
    s.write_csv(OUT / "summary.csv"); pl.Config.set_tbl_rows(100); print(s.filter(pl.col("level") == "municipality")); print("done")


if __name__ == "__main__":
    main()