#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task A — pseudo-real-time property-level forecasting. Quarterly origins from FIRST_ORIGIN; at each origin every model is
estimated on sales ≤ origin (features at their own month) and forecasts the log price of every sale in the next 12 months
(features at the origin). Predictions are stored per (origin, model, tx_id, h); scoring happens in 30_evaluate.py.
Usage: python3 scripts/10_task_a.py [--fast] [--origins 2022-12,2023-03,...]"""
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 import data as D
from wp23.models_a import TAUS, registry

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


def origins_quarterly(first: str, last_origin: str) -> list[str]:
    o = first; out = []
    while o <= last_origin:
        out.append(o); o = D.month_add(o, 3)
    return out


def main() -> None:
    ap = argparse.ArgumentParser(); ap.add_argument("--fast", action="store_true"); ap.add_argument("--family", type=str, default="trees"); ap.add_argument("--origins", type=str, default="")
    a = ap.parse_args()
    sales = D.load_sales()
    origins = a.origins.split(",") if a.origins else origins_quarterly(C.FIRST_ORIGIN, D.month_add(C.LAST_MONTH, -12))
    rows, timing = [], []
    for o in origins:
        tr, te = D.task_a_sets(sales, o)
        Xtr, ytr, cols, cats = D.design(tr); Xte, yte, _, _ = D.design(te, cats)
        meta_tr = {"mun_code": tr["mun_code"].to_list(), "mrc_code": tr["mrc_code"].to_list(), "region_code": tr["region_code"].to_list(), "h": tr["h"].to_numpy()}
        meta_te = {"mun_code": te["mun_code"].to_list(), "mrc_code": te["mrc_code"].to_list(), "region_code": te["region_code"].to_list(), "h": te["h"].to_numpy()}
        # naive benchmark: municipal median log price of the last 3 months at the origin, plus the class/subtype hedonic adjustment from OLS on the last 12 months
        base = pl.DataFrame({"tx_id": te["tx_id"], "h": te["h"], "y": yte, "naive_mun_median": te["mun_medlp_3"].fill_null(te["reg_q_3"]).to_numpy()})
        for m in registry(a.family, a.fast):
            t0 = time.time()
            try:
                m.fit(Xtr, ytr, meta_tr); p = m.predict(Xte, meta_te)
                Q = m.predict_quantiles(Xte, meta_te) if hasattr(m, "predict_quantiles") else None
            except Exception as ex:
                print("FAILED", o, m.name, repr(ex), flush=True); continue
            d = base.select(["tx_id", "h"]).with_columns(pl.lit(o).alias("origin"), pl.lit(m.name).alias("model"), pl.Series("pred", np.asarray(p, dtype=np.float64)))
            if Q is not None:
                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({"origin": o, "model": m.name, "seconds": round(time.time() - t0, 1), "n_train": len(ytr), "n_test": len(yte)})
            print(o, m.name, f"{time.time() - t0:.0f}s", f"MAE={np.mean(np.abs(p - yte)):.4f}", flush=True)
        if a.family == "trees":
            rows.append(base.select(["tx_id", "h"]).with_columns(pl.lit(o).alias("origin"), pl.lit("naive_mun_median").alias("model"), pl.Series("pred", base["naive_mun_median"])))
            base.select(["tx_id", "h", "y"]).with_columns(pl.lit(o).alias("origin")).write_parquet(OUT / f"targets_{o}.parquet")
        pl.concat(rows, how="diagonal").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()