#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 6 — validation experiments on the real data (protocol §13–15). Sub-commands (each writes results/validation/*):
  repeat      municipal/regional repeat-sales agreement (SHARP vs Case–Shiller repeat sales where ≥ 300 pairs)
  composition synthetic quality-composition shock (drop the smallest 40 % of houses from 2024-01): median vs TD vs SHARP
  placebo     months permuted within municipality × class: estimated annual changes and CI false-positive rate
  moran       residual spatial autocorrelation (Moran's I on k-nearest neighbours) of stage-1 residuals by month
  bootstrap   cluster bootstrap (by fine-location cell) of SHARP condo indices vs model-based intervals
  aggregation stock- / transaction- / equal-weighted bottom-up provincial aggregates vs the direct provincial path
  ablation    real-data ablations: variants vs full model (path differences, out-of-time innovation error), border sensitivity
  trim        sensitivity of the provincial path to the assessment-ratio screen, 1/99 % price trims, type conflicts, new construction
"""
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 sharp_hpi import config as C
from sharp_hpi import estimators as E
from sharp_hpi.hedonic import assign_loc_fe, build_features
from sharp_hpi.sharp import LEVELS, cell_log_index, run_sharp, stage1

OUT = C.RESULTS / "validation"
OUT.mkdir(parents=True, exist_ok=True)
GRID = E.month_grid("2021-01", "2026-10")


def load():
    return build_features(pl.read_parquet(C.ANALYTIC_TX))


def panel_path(level, gid, cls):
    p = pl.read_parquet(C.PROCESSED / "index_panel.parquet").filter((pl.col("geography_level") == level) & (pl.col("geography_id") == gid) & (pl.col("property_type") == cls)).sort("month")
    return np.log(p["index_value"].to_numpy() / 100.0)


def cmd_repeat(f):
    rows = []
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        for level, key in [("province", None), ("region", "region_code"), ("municipality", "mun_code")]:
            ids = ["QC"] if key is None else d.group_by(key).agg(pl.len().alias("n")).filter(pl.col("n") >= 2000)[key].to_list()
            for gid in ids:
                dc = d if key is None else d.filter(pl.col(key) == gid)
                rs = E.repeat_sales(dc, GRID, weighted=True)
                if rs.meta.get("pairs", 0) < 300:
                    continue
                s = panel_path(level, gid, cls)
                ok = ~np.isnan(rs.log_index) & ~np.isnan(s)
                a = rs.log_index[ok] - np.mean(rs.log_index[ok]); b = s[ok] - np.mean(s[ok])
                da, db = np.diff(rs.log_index), np.diff(s); okd = ~np.isnan(da) & ~np.isnan(db)
                rows.append({"class": cls, "level": level, "geo_id": gid, "pairs": rs.meta["pairs"], "months": int(ok.sum()), "rmse_level": float(np.sqrt(np.mean((a - b) ** 2))),
                             "corr_level": float(np.corrcoef(a, b)[0, 1]), "total_change_rs": float(rs.log_index[ok][-1] - rs.log_index[ok][0]), "total_change_sharp": float(s[ok][-1] - s[ok][0]),
                             "corr_12m_changes": float(np.corrcoef((rs.log_index[12:] - rs.log_index[:-12])[~np.isnan(rs.log_index[12:] - rs.log_index[:-12]) & ~np.isnan(s[12:] - s[:-12])], (s[12:] - s[:-12])[~np.isnan(rs.log_index[12:] - rs.log_index[:-12]) & ~np.isnan(s[12:] - s[:-12])])[0, 1])})
    pl.DataFrame(rows).write_csv(OUT / "repeat_sales_agreement.csv"); print(pl.DataFrame(rows))


def cmd_composition(f):
    rows = []
    for reg in ["06", "03", "07", "16"]:
        d = assign_loc_fe(f.filter((pl.col("property_class") == "house") & (pl.col("region_code") == reg)), min_n=15)
        q40 = d["fa_f"].quantile(0.4)
        shocked = d.filter(~((pl.col("month") >= "2024-01") & (pl.col("fa_f") < q40)))
        for nm, dd in [("original", d), ("shock_drop_small_from_2024", shocked)]:
            med = E.median_index(dd, GRID).log_index; td = E.pooled_td(dd, GRID).log_index
            r = run_sharp(dd, GRID); sh = r.stage1.mu
            for meth, li in [("A_median", med), ("C_pooled_td", td), ("S_sharp", sh)]:
                pre = np.nanmean(li[[GRID.index(m) for m in GRID if "2023-01" <= m <= "2023-12"]]); post = np.nanmean(li[[GRID.index(m) for m in GRID if "2024-01" <= m <= "2024-12"]])
                rows.append({"region": reg, "sample": nm, "method": meth, "n": dd.height, "change_2023_to_2024": float(post - pre)})
    out = pl.DataFrame(rows)
    out = out.pivot(on="sample", index=["region", "method"], values="change_2023_to_2024").with_columns((pl.col("shock_drop_small_from_2024") - pl.col("original")).alias("shock_effect"))
    out.write_csv(OUT / "composition_shock.csv"); print(out)


def cmd_placebo(f):
    rng = np.random.default_rng(C.SEED)
    rows = []
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        for rep in range(3):
            # permute months within municipality (keeps composition, location and overall monthly volumes within cell)
            perm = d.with_columns(pl.Series("_r", rng.random(d.height))).sort(["mun_code", "_r"])
            months_sorted = d.sort(["mun_code", "date"])["month"].to_numpy()
            perm = perm.with_columns(pl.Series("month", months_sorted)).drop("_r")
            # but keep the province-wide path: re-attach the true month effect? No — the placebo tests the *deviation* machinery:
            r = run_sharp(perm, GRID)
            for (lvl, gid), c in r.cells.items():
                dev = c.fit.smoothed; v = c.fit.smoothed_var
                d12 = dev[12:] - dev[:-12]; s12 = np.sqrt(v[12:] + v[:-12])
                rows.append({"class": cls, "rep": rep, "level": lvl, "geo_id": gid, "n_total": int(c.n.sum()), "sd_dev": float(np.std(dev)), "max_abs_12m_change": float(np.max(np.abs(d12))),
                             "share_12m_outside_95": float(np.mean(np.abs(d12) > 1.96 * s12)), "tau": float(np.sqrt(c.fit.tau2))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "placebo_cells.csv")
    summ = out.group_by(["class", "level"]).agg(pl.len().alias("cells"), pl.col("sd_dev").median(), pl.col("max_abs_12m_change").median(), pl.col("share_12m_outside_95").mean(), pl.col("tau").median()).sort(["class", "level"])
    summ.write_csv(OUT / "placebo_summary.csv"); print(summ)


def cmd_moran(f):
    from scipy.spatial import cKDTree
    rows = []
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        s1 = stage1(d, GRID)
        e = s1.resid.join(d.select(["tx_id", "x_m", "y_m"]), on="tx_id").filter(pl.col("e").is_not_null() & pl.col("e").is_not_nan())
        for scope, sub in [("province", e), ("Montréal", e.filter(pl.col("mun_code") == "66023")), ("Québec", e.filter(pl.col("mun_code") == "23027")), ("Gatineau", e.filter(pl.col("mun_code") == "81017"))]:
            for mth in ["2021-06", "2022-06", "2023-06", "2024-06", "2025-06", "2026-06"]:
                m = sub.filter(pl.col("month") == mth)
                if m.height < 200:
                    continue
                xy = m.select(["x_m", "y_m"]).to_numpy(); z = m["e"].to_numpy(); z = z - z.mean()
                k = 10; tree = cKDTree(xy); dist, idx = tree.query(xy, k=k + 1); idx = idx[:, 1:]
                lag = z[idx].mean(axis=1); I = float(np.sum(z * lag) / np.sum(z * z))
                # permutation reference
                perm = [float(np.sum(z * rng_z[idx].mean(axis=1)) / np.sum(z * z)) for rng_z in [np.random.default_rng(s).permutation(z) for s in range(99)]]
                rows.append({"class": cls, "scope": scope, "month": mth, "n": m.height, "moran_I_knn10": I, "perm_mean": float(np.mean(perm)), "perm_sd": float(np.std(perm)), "z": float((I - np.mean(perm)) / np.std(perm))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "moran_residuals.csv"); print(out)


def cmd_bootstrap(f, B: int = 40):
    rng = np.random.default_rng(C.SEED)
    d = assign_loc_fe(f.filter(pl.col("property_class") == "condo"), min_n=15)
    r0 = run_sharp(d, GRID)
    targets = [("province", "QC")] + [k for k in r0.cells if k[0] == "municipality"][:12] + [k for k in r0.cells if k[0] == "region"][:5]
    base = {k: (r0.stage1.mu if k[0] == "province" else cell_log_index(r0, *k)[0]) for k in targets}
    sev = {k: (r0.stage1.mu_se if k[0] == "province" else np.sqrt(cell_log_index(r0, *k)[1])) for k in targets}
    clusters = d["loc_fe"].unique().to_list(); boots = {k: [] for k in targets}
    for b in range(B):
        pick = rng.choice(clusters, size=len(clusters), replace=True)
        cnt = pl.DataFrame({"loc_fe": pick}).group_by("loc_fe").agg(pl.len().alias("_w"))
        db = d.join(cnt, on="loc_fe", how="inner")
        db = db.select(pl.exclude("_w").repeat_by(pl.col("_w")).explode()) if False else db.with_columns(pl.col("_w")).select(pl.all().repeat_by("_w")).explode(pl.all()).drop("_w")
        rb = run_sharp(db, GRID)
        for k in targets:
            if k[0] == "province":
                boots[k].append(rb.stage1.mu)
            elif k in rb.cells:
                boots[k].append(cell_log_index(rb, *k)[0])
        print("boot", b, flush=True)
    rows = []
    for k in targets:
        arr = np.array(boots[k]); arr = arr - arr[:, :12].mean(axis=1, keepdims=True); bse = arr.std(axis=0)
        ref = base[k] - base[k][:12].mean(); se = sev[k]
        d12b = (arr[:, 12:] - arr[:, :-12]).std(axis=0); d12m = np.sqrt(sev[k][12:] ** 2 + sev[k][:-12] ** 2)
        rows.append({"level": k[0], "geo_id": k[1], "n": int(d.filter(pl.col(dict(LEVELS).get(k[0], "mun_code")) == k[1]).height) if k[0] != "province" else d.height,
                     "model_se_level_mean": float(np.mean(se[12:])), "boot_se_level_mean": float(np.mean(bse[12:])), "ratio_level": float(np.mean(bse[12:]) / np.mean(se[12:])),
                     "model_se_12m_mean": float(np.mean(d12m)), "boot_se_12m_mean": float(np.mean(d12b)), "ratio_12m": float(np.mean(d12b) / np.mean(d12m))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "bootstrap_vs_model_se.csv"); print(out)


def cmd_aggregation(f):
    panel = pl.read_parquet(C.PROCESSED / "index_panel.parquet")
    stock = pl.read_parquet(C.RAW_UNITS_STOCK).filter(pl.col("cubf") == 1000).with_columns(
        pl.when(pl.col("nb_logements") >= 6).then(pl.lit("multi")).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"))
    sw = stock.group_by(["code_mun", "cls"]).agg(pl.len().alias("stock_units"), pl.col("valeur_role").sum().alias("stock_value")).rename({"code_mun": "geography_id"})
    sw.write_csv(OUT / "stock_weights_by_municipality_class.csv")
    rows = []
    for cls in C.PROPERTY_TYPES:
        mun = panel.filter((pl.col("geography_level") == "municipality") & (pl.col("property_type") == cls)).with_columns(np.log(pl.col("index_value") / 100).alias("li"))
        prov = panel.filter((pl.col("geography_level") == "province") & (pl.col("property_type") == cls)).sort("month").with_columns(np.log(pl.col("index_value") / 100).alias("li"))
        tw = mun.group_by("geography_id").agg(pl.col("actual_transaction_count").sum().alias("tx"))
        mun = mun.join(sw.filter(pl.col("cls") == cls).select(["geography_id", "stock_units", "stock_value"]), on="geography_id", how="left").join(tw, on="geography_id").with_columns(pl.lit(1.0).alias("eq"))
        for wname in ["stock_units", "stock_value", "tx", "eq"]:
            agg = mun.filter(pl.col(wname).is_not_null()).group_by("month").agg(((pl.col("li") * pl.col(wname)).sum() / pl.col(wname).sum()).alias("li"), pl.col(wname).sum().alias("w")).sort("month")
            a = agg.join(prov.select(["month", pl.col("li").alias("li_prov")]), on="month")
            li = a["li"].to_numpy().copy(); lp = a["li_prov"].to_numpy().copy(); li -= li[:12].mean(); lp -= lp[:12].mean()
            rows.append({"class": cls, "weights": wname, "cells": mun.filter(pl.col(wname).is_not_null())["geography_id"].n_unique(), "rmse_vs_direct": float(np.sqrt(np.mean((li - lp) ** 2))),
                         "total_change_bottomup": float(li[-1] - li[0]), "total_change_direct": float(lp[-1] - lp[0]), "max_abs_gap": float(np.max(np.abs(li - lp)))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "aggregation_sensitivity.csv"); print(out)


def cmd_ablation(f):
    rows = []
    variants = {"full": {}, "noEB": {"eb": False}, "noTV": {"time_varying": False}, "LLT": {"model": "llt"}, "norobust": {"robust": False},
                "coarseloc": {"min_loc_n": 10**9}, "skipMRC": {"levels": [("region", "region_code"), ("municipality", "mun_code")]}}
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        runs = {}
        for v, kw in variants.items():
            t0 = time.time(); runs[v] = run_sharp(d, GRID, **kw); print(cls, v, f"{time.time() - t0:.0f}s", flush=True)
        full = runs["full"]
        for v, r in runs.items():
            for lvl in ["region", "mrc", "municipality"]:
                keys = [k for k in r.cells if k[0] == lvl and k in full.cells]
                diffs, oos, cov, gains = [], [], [], []
                for k in keys:
                    p, var = cell_log_index(r, *k); p0, _ = cell_log_index(full, *k)
                    diffs.append(np.sqrt(np.mean((np.diff(p) - np.diff(p0)) ** 2)))
                    c = r.cells[k]; z = c.oos_z[np.isfinite(c.oos_z)]
                    if len(z) > 5:
                        oos.append(np.sqrt(np.mean(z ** 2))); cov.append(np.mean(np.abs(z) < 1.96))
                    gains.append(np.nanmean(c.fit.gain[c.n > 0]) if (c.n > 0).any() else np.nan)
                rows.append({"class": cls, "variant": v, "level": lvl, "cells": len(keys), "rmse_monthly_change_vs_full": float(np.mean(diffs)) if diffs else np.nan,
                             "oos_innovation_rmse_z": float(np.mean(oos)) if oos else np.nan, "oos_innovation_share_within_95": float(np.mean(cov)) if cov else np.nan,
                             "mean_gain": float(np.nanmean(gains))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "ablation_real.csv"); print(out)


def cmd_trim(f):
    clean = pl.read_parquet(C.CLEAN_TX)
    from sharp_hpi.cleaning import analytic
    rows = []
    for cls in ["house", "condo"]:
        base = f.filter(pl.col("property_class") == cls)
        variants = {"headline": base}
        # include assessment-ratio outliers back
        inc = clean.with_columns(pl.when(pl.col("exclude_reason_headline") == "assessment_ratio_outlier").then(None).otherwise(pl.col("exclude_reason_headline")).alias("exclude_reason_headline")).with_columns(pl.col("exclude_reason_headline").is_null().alias("in_headline_sample"))
        variants["with_ratio_outliers"] = build_features(analytic(inc)).filter(pl.col("property_class") == cls)
        variants["drop_type_conflicts"] = base.filter(~pl.col("type_conflict"))
        variants["drop_new_construction"] = base.filter(~pl.col("new_construction"))
        q = base.group_by("month").agg(pl.col("price").quantile(0.01).alias("q1"), pl.col("price").quantile(0.99).alias("q99"))
        variants["trim_1_99_monthly"] = base.join(q, on="month").filter((pl.col("price") > pl.col("q1")) & (pl.col("price") < pl.col("q99")))
        variants["provider_only_classes"] = base.filter(pl.col("class_source") == "roll")
        ref = None
        for nm, dd in variants.items():
            s = E.pooled_td(assign_loc_fe(dd, min_n=15), GRID, robust=True).log_index; s_ols = E.pooled_td(assign_loc_fe(dd, min_n=15), GRID, robust=False).log_index
            if ref is None:
                ref = s
            rows.append({"class": cls, "variant": nm, "n": dd.height, "total_change_huber": float(s[-1] - s[0]), "total_change_ols": float(s_ols[-1] - s_ols[0]),
                         "rmse_monthly_vs_headline": float(np.sqrt(np.nanmean((np.diff(s) - np.diff(ref)) ** 2))), "max_abs_level_gap": float(np.nanmax(np.abs(s - ref)))})
    out = pl.DataFrame(rows); out.write_csv(OUT / "trimming_sensitivity.csv"); print(out)


if __name__ == "__main__":
    cmd = sys.argv[1]
    f = load()
    {"repeat": cmd_repeat, "composition": cmd_composition, "placebo": cmd_placebo, "moran": cmd_moran, "bootstrap": cmd_bootstrap,
     "aggregation": cmd_aggregation, "ablation": cmd_ablation, "trim": cmd_trim}[cmd](f)