#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Additional analyses (protocol §8–§11): temporal evolution of errors, regime-dependent accuracy (ex-post regimes and the
RAREF origin-dated regime probabilities), directional accuracy and turning points, probabilities of decline (Brier score,
reliability), Holm-adjusted multiple-testing, equal-weight combinations, municipal performance distributions, computational
cost, economic magnitude, and transaction-selection rates. Writes results/extra/*.csv."""
from __future__ import annotations

import glob
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 scipy import stats

from wp23 import config as C
from wp23 import evaluation as E
from wp23.data import month_add

OUT = C.RESULTS / "extra"; OUT.mkdir(parents=True, exist_ok=True)
TAUS = np.array([0.05, 0.25, 0.5, 0.75, 0.95])
KEY_B = ["no_change", "mean_growth", "momentum", "ar2_cell", "ridge_panel", "lightgbm_panel", "tft_lite", "bhpf", "raref", "qanss", "dsgfn", "qrefa"]


def load_b():
    p = pl.concat([pl.read_parquet(f) for f in sorted(glob.glob(str(C.RESULTS / "task_b" / "predictions_*.parquet")))], how="diagonal")
    return p.with_columns((pl.col("pred") - pl.col("target")).abs().alias("ae"), ((pl.col("pred") - pl.col("target")) ** 2).alias("se"))


def prob_below(Q: np.ndarray, x: float = 0.0) -> np.ndarray:
    """P(target < x) from the 5-quantile forecast by piecewise-linear interpolation of the quantile function (tails extrapolated with the 5–25 and 75–95 slopes)."""
    out = np.empty(len(Q))
    for i, q in enumerate(Q):
        q = np.sort(q); lo, hi = q[0] - (q[1] - q[0]) * 2, q[-1] + (q[-1] - q[-2]) * 2
        xs = np.concatenate([[lo], q, [hi]]); ps = np.concatenate([[0.0], TAUS, [1.0]])
        out[i] = float(np.interp(x, xs, ps)) if np.all(np.isfinite(xs)) else np.nan
    return out


def temporal_b(p):
    d = p.filter(pl.col("model").is_in(KEY_B)).group_by(["model", "h", "origin"]).agg(pl.col("ae").mean().alias("mae"), pl.len().alias("n")).sort(["model", "h", "origin"])
    d.write_csv(OUT / "error_by_origin_b.csv")
    y = p.filter(pl.col("model").is_in(KEY_B)).with_columns(pl.col("origin").str.slice(0, 4).alias("year")).group_by(["model", "h", "year"]).agg(pl.col("ae").mean().alias("mae"), pl.len().alias("n"))
    nc = y.filter(pl.col("model") == "no_change").select(["h", "year", pl.col("mae").alias("mae_nc")])
    y.join(nc, on=["h", "year"]).with_columns((pl.col("mae") / pl.col("mae_nc")).alias("mase")).sort(["h", "year", "mae"]).write_csv(OUT / "error_by_year_b.csv")


def regimes_b(p):
    # ex-post regimes from the realised provincial 3-month house change at the origin's horizon window: rising (> +2 pts), falling (< −2 pts), flat
    prov = p.filter((pl.col("level") == "province") & (pl.col("class") == "house") & (pl.col("model") == "no_change")).select(["origin", "h", pl.col("target").alias("prov_target")])
    d = p.filter(pl.col("model").is_in(KEY_B)).join(prov, on=["origin", "h"], how="left").with_columns(
        pl.when(pl.col("prov_target") > 0.02).then(pl.lit("rising")).when(pl.col("prov_target") < -0.02).then(pl.lit("falling")).otherwise(pl.lit("flat")).alias("regime"))
    g = d.group_by(["model", "h", "regime"]).agg(pl.col("ae").mean().alias("mae"), pl.len().alias("n"), pl.col("origin").n_unique().alias("origins"))
    nc = g.filter(pl.col("model") == "no_change").select(["h", "regime", pl.col("mae").alias("mae_nc")])
    g.join(nc, on=["h", "regime"]).with_columns((pl.col("mae") / pl.col("mae_nc")).alias("mase")).sort(["h", "regime", "mae"]).write_csv(OUT / "error_by_regime_b.csv")
    # RAREF regime probabilities by origin (origin-dated): refit the HMM gating on the h=3 panel with expanding training
    from wp23.models_b import RAREF
    panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet").filter(pl.col("h") == 3)
    origins = sorted(panel["origin"].unique().to_list()); rows = []
    for o in origins[10:]:
        tr = panel.filter(pl.col("origin").map_elements(lambda x: month_add(x, 3), return_dtype=pl.Utf8) <= o); te = panel.filter(pl.col("origin") == o)
        if tr["origin"].n_unique() < 10: continue
        m = RAREF().fit(tr); both = pl.concat([tr.select(te.columns), te], how="diagonal"); P = m._probs(both)[-te.height:]
        for (cls, lvl), pr in zip(te.select(["class", "level"]).iter_rows(), P):
            if lvl == "province":
                hmm = m.hmm_for(m._growth_series(both, o, cls)); order = np.argsort(hmm.mu)  # states ordered by mean growth
                rows.append({"origin": o, "class": cls, "p_low": float(pr[order[0]]), "p_mid": float(pr[order[1]]), "p_high": float(pr[order[2]]), "mu_low": float(hmm.mu[order[0]]), "mu_high": float(hmm.mu[order[2]]),
                             "prov_growth_last": float(np.diff(np.array(te.filter((pl.col("class") == cls) & (pl.col("level") == "province"))["hist"][0]))[-1])})
    pl.DataFrame(rows).write_csv(OUT / "raref_regimes.csv")


def direction_b(p):
    d = p.filter(pl.col("model").is_in(KEY_B) & (pl.col("target").abs() > 0.005)).with_columns((np.sign(pl.col("pred")) == np.sign(pl.col("target"))).cast(pl.Float64).alias("hit"))
    g = d.group_by(["model", "h"]).agg(pl.col("hit").mean().alias("directional_accuracy"), pl.len().alias("n"))
    # turning points: origins where the realised change of the cell changes sign relative to the previous origin's realised change
    s = p.filter(pl.col("model") == "no_change").sort(["class", "level", "geo_id", "h", "origin"]).with_columns(pl.col("target").shift(1).over(["class", "level", "geo_id", "h"]).alias("prev"))
    tp = s.filter((np.sign(pl.col("target")) != np.sign(pl.col("prev"))) & pl.col("prev").is_not_null() & (pl.col("target").abs() > 0.01)).select(["class", "level", "geo_id", "h", "origin"]).with_columns(pl.lit(True).alias("turning"))
    dd = d.join(tp, on=["class", "level", "geo_id", "h", "origin"], how="left").with_columns(pl.col("turning").fill_null(False))
    g2 = dd.filter(pl.col("turning")).group_by(["model", "h"]).agg(pl.col("hit").mean().alias("directional_accuracy_at_turning_points"), pl.len().alias("n_turning"))
    g.join(g2, on=["model", "h"], how="left").sort(["h", "directional_accuracy"], descending=[False, True]).write_csv(OUT / "directional_b.csv")


def decline_b(p):
    rows, rel = [], []
    for (m, h), d in p.filter(pl.col("model").is_in(KEY_B)).group_by(["model", "h"]):
        Q = d.select(["q05", "q25", "q50", "q75", "q95"]).to_numpy(); y = d["target"].to_numpy(); ok = np.all(np.isfinite(Q), axis=1)
        pd_ = prob_below(Q[ok], 0.0); yd = (y[ok] < 0).astype(float); pa = 1 - prob_below(Q[ok], 0.02); ya = (y[ok] > 0.02).astype(float)
        brier_d = float(np.mean((pd_ - yd) ** 2)); brier_a = float(np.mean((pa - ya) ** 2)); base_d = float(np.mean((yd.mean() - yd) ** 2)); base_a = float(np.mean((ya.mean() - ya) ** 2))
        rows.append({"model": m, "h": h, "n": int(ok.sum()), "share_declines": float(yd.mean()), "brier_decline": brier_d, "brier_skill_decline": 1 - brier_d / base_d, "brier_up2": brier_a, "brier_skill_up2": 1 - brier_a / base_a})
        bins = np.linspace(0, 1, 6)
        for b in range(5):
            sel = (pd_ >= bins[b]) & (pd_ < bins[b + 1] + (b == 4))
            if sel.sum() >= 30: rel.append({"model": m, "h": h, "bin": b, "p_mean": float(pd_[sel].mean()), "freq": float(yd[sel].mean()), "n": int(sel.sum())})
    pl.DataFrame(rows).sort(["h", "brier_decline"]).write_csv(OUT / "decline_b.csv"); pl.DataFrame(rel).write_csv(OUT / "decline_reliability_b.csv")


def holm_tests():
    for tag in ["a", "b"]:
        f = C.RESULTS / f"eval_{tag}" / "tests_vs_reference.csv"
        if not f.exists(): continue
        t = pl.read_csv(f, infer_schema_length=None); out = []
        for h, d in t.group_by("h"):
            d = d.filter(pl.col("dm_p").is_not_null()).sort("dm_p"); pv = d["dm_p"].to_numpy(); m = len(pv); adj = np.minimum.accumulate(np.minimum(1, pv * (m - np.arange(m)))[::-1])[::-1]
            adj = np.maximum.accumulate(np.minimum(1, pv * (m - np.arange(m))))  # Holm step-down
            out.append(d.with_columns(pl.Series("dm_p_holm", adj), pl.Series("dm_p_bonferroni", np.minimum(1, pv * m))))
        pl.concat(out).write_csv(OUT / f"tests_holm_{tag}.csv")


def combos_b(p):
    base = p.filter(pl.col("model").is_in(["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell", "no_change", "qrefa"])).select(["class", "level", "geo_id", "origin", "h", "target", "model", "pred"])
    w = base.pivot(on="model", index=["class", "level", "geo_id", "origin", "h", "target"], values="pred").drop_nulls()
    combos = {"eq_lgbm_qanss": ["lightgbm_panel", "qanss"], "eq_lgbm_qanss_ridge_ar2": ["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell"], "eq_lgbm_ar2": ["lightgbm_panel", "ar2_cell"], "eq_all_with_nochange": ["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell", "no_change"]}
    rows = []
    for h, d in w.group_by("h"):
        h = h[0]; y = d["target"].to_numpy(); nc = np.abs(d["no_change"].to_numpy() - y).mean()
        for name, cols in combos.items():
            pr = np.mean(np.column_stack([d[c].to_numpy() for c in cols]), axis=1); rows.append({"h": h, "combination": name, "mae": float(np.mean(np.abs(pr - y))), "mase": float(np.mean(np.abs(pr - y)) / nc), "n": d.height})
        for c in ["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell", "qrefa"]:
            rows.append({"h": h, "combination": c, "mae": float(np.mean(np.abs(d[c].to_numpy() - y))), "mase": float(np.mean(np.abs(d[c].to_numpy() - y)) / nc), "n": d.height})
    pl.DataFrame(rows).sort(["h", "mae"]).write_csv(OUT / "combinations_b.csv")


def municipal_distribution_b(p):
    d = p.filter((pl.col("level") == "municipality") & (pl.col("h") == 3) & pl.col("model").is_in(["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell", "no_change"]))
    g = d.group_by(["class", "geo_id", "model"]).agg(pl.col("ae").mean().alias("mae"), pl.col("n_last").mean().alias("mean_n"), pl.len().alias("n"))
    nc = g.filter(pl.col("model") == "no_change").select(["class", "geo_id", pl.col("mae").alias("mae_nc")])
    g = g.join(nc, on=["class", "geo_id"]).with_columns((pl.col("mae") / pl.col("mae_nc")).alias("mase")); g.write_csv(OUT / "municipal_mase_b.csv")
    g.filter(pl.col("model") != "no_change").group_by("model").agg(pl.col("mase").median().alias("median_mase"), (pl.col("mase") < 1).mean().alias("share_beating_no_change"), pl.col("mase").quantile(0.1).alias("p10"), pl.col("mase").quantile(0.9).alias("p90"), pl.len().alias("cells")).sort("median_mase").write_csv(OUT / "municipal_mase_summary_b.csv")


def cost():
    rows = []
    for tag in ["a", "b"]:
        t = pl.concat([pl.read_csv(f) for f in sorted(glob.glob(str(C.RESULTS / f"task_{tag}" / "timing_*.csv")))])
        g = t.group_by("model").agg(pl.col("seconds").mean().alias("mean_seconds_per_fit"), pl.col("seconds").sum().alias("total_seconds"), pl.len().alias("fits"), pl.col("n_train").mean().alias("mean_train_rows")).with_columns(pl.lit(tag.upper()).alias("task"))
        rows.append(g)
    pl.concat(rows).sort(["task", "mean_seconds_per_fit"]).write_csv(OUT / "computational_cost.csv")


def temporal_a():
    files = sorted(glob.glob(str(C.RESULTS / "task_a" / "predictions_*.parquet")))
    if not files: return
    pred = pl.concat([pl.read_parquet(f) for f in files], how="diagonal").filter(pl.col("model").is_in(["lightgbm", "hshn_partial", "hedonic_ols", "naive_mun_median", "xgboost", "random_forest", "mlp"]))
    tg = pl.concat([pl.read_parquet(f) for f in sorted(glob.glob(str(C.RESULTS / "task_a" / "targets_*.parquet")))])
    j = pred.join(tg, on=["tx_id", "h", "origin"]).with_columns((pl.col("pred") - pl.col("y")).abs().alias("ae"), (pl.col("pred") - pl.col("y")).alias("err"))
    j.group_by(["model", "origin"]).agg(pl.col("ae").mean().alias("mae"), pl.col("err").mean().alias("bias"), pl.len().alias("n")).sort(["model", "origin"]).write_csv(OUT / "error_by_origin_a.csv")
    j.filter(pl.col("model").is_in(["lightgbm", "hshn_partial", "hedonic_ols"])).group_by(["model", "h"]).agg(pl.col("ae").mean().alias("mae"), pl.col("err").mean().alias("bias")).sort(["model", "h"]).write_csv(OUT / "error_by_month_horizon_a.csv")
    # economic magnitude: MAE in dollars at the class median price
    sales = pl.read_parquet(C.WP22_ANALYTIC, columns=["tx_id", "property_class", "price"]); jj = j.join(sales, on="tx_id")
    jj.filter(pl.col("model").is_in(["lightgbm", "hedonic_ols", "naive_mun_median"]) & (pl.col("h") <= 3)).group_by(["model", "property_class"]).agg(pl.col("ae").mean().alias("mae_log"), pl.col("price").median().alias("median_price"), (pl.col("pred").exp() - pl.col("price")).abs().median().alias("median_abs_dollar_error")).sort(["property_class", "mae_log"]).write_csv(OUT / "dollar_errors_a.csv")


def selection():
    """Transaction selection: annual sale rate of the roll stock by class, floor-area and age decile (sold units / stock units)."""
    st = pl.read_parquet(C.WP22_STOCK, columns=["id_provinc", "cubf", "nb_logements", "lien_physique", "annee_construction", "aire_etages_m2"]).filter(pl.col("cubf") == 1000)
    st = st.with_columns(pl.when(pl.col("nb_logements") >= 6).then(None).when(pl.col("nb_logements").is_between(2, 5)).then(pl.lit("plex")).when((pl.col("nb_logements") <= 1) & (pl.col("lien_physique") == "Intégré")).then(pl.lit("condo")).when(pl.col("nb_logements") <= 1).then(pl.lit("house")).otherwise(None).alias("cls")).filter(pl.col("cls").is_not_null())
    sold = pl.read_parquet(C.WP22_ANALYTIC, columns=["parcel_id", "property_class"]).group_by(["parcel_id", "property_class"]).agg(pl.len().alias("sales"))
    j = st.join(sold, left_on=["id_provinc", "cls"], right_on=["parcel_id", "property_class"], how="left").with_columns(pl.col("sales").fill_null(0))
    rows = []
    for cls in ["house", "condo", "plex"]:
        d = j.filter(pl.col("cls") == cls).with_columns(pl.col("aire_etages_m2").qcut(5, labels=["fa1", "fa2", "fa3", "fa4", "fa5"]).alias("fa_q"), (2026 - pl.col("annee_construction")).qcut(5, labels=["age1", "age2", "age3", "age4", "age5"]).alias("age_q"))
        for col in ["fa_q", "age_q"]:
            for g, s in d.group_by(col):
                rows.append({"class": cls, "dimension": col, "quintile": g[0], "stock_units": s.height, "sales_2021_2026": int(s["sales"].sum()), "annual_sale_rate_pct": 100 * float(s["sales"].sum()) / s.height / 5.75})
    pl.DataFrame(rows).sort(["class", "dimension", "quintile"]).write_csv(OUT / "selection_rates.csv")


if __name__ == "__main__":
    p = load_b()
    for fn in [lambda: temporal_b(p), lambda: regimes_b(p), lambda: direction_b(p), lambda: decline_b(p), holm_tests, lambda: combos_b(p), lambda: municipal_distribution_b(p), cost, temporal_a, selection]:
        try: fn(); print("ok", getattr(fn, "__name__", "λ"))
        except Exception as ex:
            import traceback; traceback.print_exc()
    print("done")