#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task B — pseudo-real-time index forecasting on the cell × origin panel.
For each horizon h and each quarterly re-estimation origin o: training rows are panel rows whose target was observable at o
(origin' + h ≤ o); models are fitted once and used for the forecasts issued at origins o, o+1, o+2 (features at those origins).
Usage: python3 scripts/20_task_b.py [--fast] [--models a,b] [--refit-every 3]"""
from __future__ import annotations

import argparse
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 TAUS, registry

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


def main() -> None:
    ap = argparse.ArgumentParser(); ap.add_argument("--fast", action="store_true"); ap.add_argument("--family", type=str, default="trees"); ap.add_argument("--models", type=str, default=""); ap.add_argument("--refit-every", type=int, default=3)
    ap.add_argument("--min-train-origins", type=int, default=10)
    a = ap.parse_args()
    panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet")
    models = registry(a.family, a.fast)
    if a.models:
        keep = set(a.models.split(",")); models = [m for m in models if m.name in keep]
    rows, timing = [], []
    for h in C.HORIZONS:
        ph = panel.filter(pl.col("h") == h); origins = sorted(ph["origin"].unique().to_list())
        for k in range(0, len(origins), a.refit_every):
            o = origins[k]; test_origins = origins[k:k + a.refit_every]
            tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)
            if tr["origin"].n_unique() < a.min_train_origins:
                continue
            te = ph.filter(pl.col("origin").is_in(test_origins))
            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", h, o, m.name, repr(ex), flush=True); continue
                d = te.select(["class", "level", "geo_id", "origin", "h", "target", "n_last", "grade_target", "region_code"]).with_columns(pl.lit(m.name).alias("model"), pl.lit(o).alias("fit_origin"), pl.Series("pred", np.asarray(p, dtype=np.float64)))
                for j, t in enumerate(TAUS):
                    d = d.with_columns(pl.Series(f"q{int(t*100):02d}", np.asarray(Q[:, j], dtype=np.float64)))
                rows.append(d); timing.append({"h": h, "fit_origin": o, "model": m.name, "seconds": round(time.time() - t0, 1), "n_train": tr.height, "n_test": te.height})
                print(h, o, m.name, f"{time.time() - t0:.1f}s", f"MAE={np.mean(np.abs(p - te['target'].to_numpy())):.4f}", flush=True)
            pl.concat(rows).write_parquet(OUT / f"predictions_{a.family}.parquet"); pl.DataFrame(timing).write_csv(OUT / f"timing_{a.family}.csv")
    print("done")


if __name__ == "__main__":
    main()