#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Summarise Task B ablations: MAE / CRPS-proxy by variant × horizon × level, with per-origin bootstrap CIs of the difference to the full model."""
from __future__ import annotations

import sys
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 evaluation as E

OUT = C.RESULTS / "ablation_b"
p = pl.read_parquet(OUT / "predictions_ablation.parquet").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"))
ref = {"dsgfn_no_distance": "dsgfn", "dsgfn_no_comovement": "dsgfn", "dsgfn_no_graph": "dsgfn", "qanss_no_parent": "qanss", "equal_weight_combo": "qrefa", "qrefa_without_qanss": "qrefa", "qrefa_without_bhpf": "qrefa", "qrefa_without_dsgfn": "qrefa", "qrefa_without_raref": "qrefa",
       "ridge_no_parent": "ridge_panel", "ridge_no_volume": "ridge_panel", "ridge_no_momentum": "ridge_panel"}
rows = []
for (m, h, lvl), d in p.group_by(["model", "h", "level"]):
    r = {"model": m, "h": h, "level": lvl, "n": d.height, "mae": float(d["ae"].mean()), "cov90": float(d["cov90"].mean())}
    base = ref.get(m)
    if base:
        b = p.filter((pl.col("model") == base) & (pl.col("h") == h) & (pl.col("level") == lvl)).select(["class", "geo_id", "origin", pl.col("ae").alias("ae_base")])
        j = d.join(b, on=["class", "geo_id", "origin"]); per = j.group_by("origin").agg(pl.col("ae").mean().alias("a"), pl.col("ae_base").mean().alias("b")).sort("origin")
        diff = (per["a"] - per["b"]).to_numpy(); r["mae_diff_vs_ref"] = float(np.mean(diff)); lo, hi = E.block_bootstrap_ci(diff, B=500, mean_block=max(2, int(h))); r["diff_lo"], r["diff_hi"] = lo, hi; r["reference"] = base
    rows.append(r)
s = pl.DataFrame(rows, infer_schema_length=None).sort(["level", "h", "mae"]); s.write_csv(OUT / "ablation_summary.csv")
pl.Config.set_tbl_rows(80); print(s.filter(pl.col("level") == "municipality").select(["model", "h", "mae", "cov90", "mae_diff_vs_ref", "diff_lo", "diff_hi"]))
if (OUT / "dsgfn_attention.csv").exists():
    a = pl.read_csv(OUT / "dsgfn_attention.csv", infer_schema_length=None)
    geo = pl.read_parquet(C.WP22_GEO, columns=["mun_code", "mrc_code", "region_code", "x_m", "y_m"]).group_by("mun_code").agg(pl.col("mrc_code").first(), pl.col("region_code").first(), pl.col("x_m").mean(), pl.col("y_m").mean())
    aa = a.filter(pl.col("level") == "municipality").join(geo.rename({"mun_code": "geo_id"}), on="geo_id", how="left").join(geo.rename({"mun_code": "top_neighbour", "mrc_code": "mrc_n", "region_code": "reg_n", "x_m": "xn", "y_m": "yn"}), on="top_neighbour", how="left")
    aa = aa.with_columns((((pl.col("x_m") - pl.col("xn")) ** 2 + (pl.col("y_m") - pl.col("yn")) ** 2).sqrt() / 1000).alias("dist_km"), (pl.col("region_code") == pl.col("reg_n")).alias("same_region"), (pl.col("mrc_code") == pl.col("mrc_n")).alias("same_mrc"))
    summ = aa.group_by("h").agg(pl.len().alias("n"), pl.col("top_weight").mean(), pl.col("entropy").mean(), pl.col("dist_km").median().alias("median_dist_top_km"), pl.col("same_region").mean(), pl.col("same_mrc").mean())
    summ.write_csv(OUT / "dsgfn_attention_summary.csv"); print(summ)