#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 5 — complete monthly SHARP-HPI panel (retrospective research series, full sample) at province, region, MRC and
municipality level for house / condo / plex, plus coverage, reliability and representative-price tables and the
QHPI-architecture benchmark panel. Writes data/processed/index_panel.parquet|csv and results/panel/*.csv."""
from __future__ import annotations

import json
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.estimators import month_grid
from sharp_hpi.hedonic import build_features
from sharp_hpi.panel import build_rows
from sharp_hpi.sharp import LEVELS, cell_log_index, run_sharp

OUT = C.RESULTS / "panel"
OUT.mkdir(parents=True, exist_ok=True)
GRID = month_grid("2021-01", "2026-10")
VINTAGE = C.SNAPSHOT_TAKEN[:10]


def names(geo: pl.DataFrame) -> dict:
    d = {}
    for lvl, key, nm in [("region", "region_code", "region_name"), ("mrc", "mrc_code", "mrc_name"), ("municipality", "mun_code", "mun_name")]:
        for k, v in geo.select([key, nm]).unique().iter_rows():
            d[(lvl, k)] = v
    return d


def main() -> None:
    a = pl.read_parquet(C.ANALYTIC_TX)
    f = build_features(a)
    nm = names(pl.read_parquet(C.GEO_TX, columns=["region_code", "region_name", "mrc_code", "mrc_name", "mun_code", "mun_name"]))
    rows, rep, cov, meta = [], [], [], {}
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        t0 = time.time()
        r = run_sharp(d, GRID)
        meta[cls] = {"n": r.stage1.n, "sigma_robust": r.stage1.sigma2 ** 0.5, "r2_within": r.stage1.r2_within, "priors": {k: [float(v[0]), float(v[1])] for k, v in r.priors.items()},
                     "seconds": round(time.time() - t0, 1), "coef": {k: float(v) for k, v in r.stage1.coef.items() if "month" not in k}}
        n_prov = np.array([d.filter(pl.col("month") == m).height for m in GRID], float)
        rows.append(build_rows("QC", "Québec (province)", "province", cls, GRID, r.stage1.mu, r.stage1.mu_se ** 2, n_prov, n_prov, np.ones(len(GRID)),
                               "SHARP-HPI v%s (stage-1 robust TD)" % C.MODEL_VERSION, C.SAMPLE_END, VINTAGE, C.MODEL_VERSION))
        for (lvl, gid), c in r.cells.items():
            p, v = cell_log_index(r, lvl, gid, "smoothed")
            ess = np.where(c.fit.smoothed_var > 0, c.sigma2_cell / np.maximum(c.fit.smoothed_var, 1e-12), np.inf)
            ess = np.minimum(ess, 1e6)
            rows.append(build_rows(gid, nm.get((lvl, gid), gid), lvl, cls, GRID, p, v, c.n, ess, c.fit.gain, "SHARP-HPI v%s (hierarchical, smoothed)" % C.MODEL_VERSION,
                                   C.SAMPLE_END, VINTAGE, C.MODEL_VERSION))
            # reliability bundle
            cov.append({"geography_level": lvl, "geography_id": gid, "property_type": cls, "n_total": int(c.n.sum()), "months_with_sales": int((c.n > 0).sum()),
                        "median_monthly_n": float(np.median(c.n[c.n > 0])) if (c.n > 0).any() else 0.0, "tau": float(np.sqrt(c.fit.tau2)), "sigma_cell": float(np.sqrt(c.sigma2_cell)),
                        "mean_gain": float(np.nanmean(c.fit.gain[c.n > 0])) if (c.n > 0).any() else 0.0, "mean_ess": float(np.mean(ess)), "mean_overlap": float(np.nanmean(c.overlap)),
                        "acf1_dev": c.rho1, "oos_z_rmse": float(np.sqrt(np.nanmean(c.oos_z ** 2))) if np.isfinite(c.oos_z).any() else np.nan,
                        "innov_weight_min": float(np.min(c.fit.innov_weights))})
            # representative price at 2021 prices: exp(mean(log P - log index)) over the cell's sample
            sub = d.filter(pl.col(dict(LEVELS)[lvl]) == gid)
            li = dict(zip(GRID, p)); adj = np.array([li[m] for m in sub["month"].to_list()])
            base = [i for i, m in enumerate(GRID) if m.startswith("2021")]
            b2021 = float(np.mean(p[base]))
            rep.append({"geography_level": lvl, "geography_id": gid, "property_type": cls, "representative_price_2021": float(np.exp(np.mean(sub["log_price"].to_numpy() - adj + b2021))),
                        "n": sub.height})
    panel = pl.concat(rows)
    panel.write_parquet(C.PROCESSED / "index_panel.parquet"); panel.write_csv(C.PROCESSED / "index_panel.csv")
    pl.DataFrame(cov).write_csv(OUT / "reliability_cells.csv")
    pl.DataFrame(rep).write_csv(OUT / "representative_prices.csv")
    (OUT / "stage1_meta.json").write_text(json.dumps(meta, indent=2))
    # coverage table: every geography × type × month
    panel.select(["geography_level", "geography_id", "geography_name", "property_type", "month", "actual_transaction_count", "effective_sample_size",
                  "reliability_grade", "direct_or_model_assisted", "suppression_flag"]).write_csv(OUT / "coverage_cells_months.csv")
    panel.group_by(["geography_level", "property_type", "reliability_grade"]).agg(pl.len().alias("cell_months")).sort(["geography_level", "property_type", "reliability_grade"]).write_csv(OUT / "grade_distribution.csv")
    print(panel.shape, meta)


if __name__ == "__main__":
    main()