#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 2 — benchmark indices A–K and SHARP at province × class and region × class (retrospective, full sample).

Output: results/benchmarks/series.parquet (long: class, level, geo_id, method, month, log_index, se, n) + timing.csv.
GBM (J) is run at province level and for the three largest regions only (cost), see docs/decision_log.md D2.1.
"""
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.qhpi import qhpi_cell, qhpi_fullpool_residuals, qhpi_province
from sharp_hpi.sharp import cell_log_index, run_sharp

OUT = C.RESULTS / "benchmarks"
OUT.mkdir(parents=True, exist_ok=True)
GRID = E.month_grid("2021-01", "2026-10")
GBM_REGIONS = {"06", "03", "16"}


def rec(rows, timing, cls, level, gid, s: E.IndexSeries, t0):
    f = s.to_frame().with_columns(pl.lit(cls).alias("class"), pl.lit(level).alias("level"), pl.lit(gid).alias("geo_id"))
    if "se" not in f.columns:
        f = f.with_columns(pl.lit(None, dtype=pl.Float64).alias("se"))
    if "n" not in f.columns:
        f = f.with_columns(pl.lit(None, dtype=pl.Float64).alias("n"))
    rows.append(f.select(["class", "level", "geo_id", "method", "month", "log_index", "se", "n"]))
    timing.append({"class": cls, "level": level, "geo_id": gid, "method": s.method, "seconds": round(time.time() - t0, 2)})
    print(cls, level, gid, s.method, f"{time.time() - t0:.1f}s", flush=True)


def main() -> None:
    a = pl.read_parquet(C.ANALYTIC_TX)
    f = build_features(a)
    rows, timing = [], []
    for cls in ("house", "condo", "plex"):
        d = f.filter(pl.col("property_class") == cls)
        prov = assign_loc_fe(d, min_n=15).with_columns(pl.col("mrc_code").alias("loc_coarse"))
        # --- province ---
        for fn, kw in [(E.median_index, {"stat": "median"}), (E.median_index, {"stat": "mean"}), (E.stratified_median, {}),
                       (E.pooled_td, {"robust": True}), (E.pooled_td, {"robust": False}),
                       (E.rolling_td, {"window": 13, "splice": "mean"}), (E.rolling_td, {"window": 13, "splice": "movement"}), (E.rolling_td, {"window": 13, "splice": "window"}),
                       (E.rolling_td, {"window": 25, "splice": "mean"}), (E.rolling_td, {"window": 7, "splice": "mean"}),
                       (E.hedonic_imputation, {"kind": "laspeyres"}), (E.hedonic_imputation, {"kind": "paasche"}), (E.hedonic_imputation, {"kind": "fisher"}),
                       (E.hedonic_imputation, {"kind": "fisher", "chained": False}), (E.geks, {"window": 25}), (E.geks, {"window": 13}),
                       (E.repeat_sales, {"weighted": True}), (E.repeat_sales, {"weighted": False}), (E.gbm_residual_index, {})]:
            t0 = time.time(); rec(rows, timing, cls, "province", "QC", fn(prov, GRID, **kw), t0)
        t0 = time.time(); sK = qhpi_province(prov, GRID); sK.method = "K_qhpi"; rec(rows, timing, cls, "province", "QC", sK, t0)
        t0 = time.time(); r = run_sharp(d, GRID); s1 = r.stage1
        rec(rows, timing, cls, "province", "QC", E.IndexSeries(GRID, s1.mu, se=s1.mu_se, n=E._counts(d, GRID), method="S_sharp"), t0)
        t0 = time.time(); r0 = run_sharp(d, GRID, time_varying=False)
        rec(rows, timing, cls, "province", "QC", E.IndexSeries(GRID, r0.stage1.mu, se=r0.stage1.mu_se, n=E._counts(d, GRID), method="S_sharp_notv"), t0)
        # QHPI hierarchy inputs
        qres, qsig = qhpi_fullpool_residuals(prov, GRID)
        # --- regions ---
        for reg in sorted(d["region_code"].unique().to_list()):
            dr = d.filter(pl.col("region_code") == reg)
            if dr.height < 500:
                continue
            drr = assign_loc_fe(dr, min_n=15).with_columns(pl.col("mun_code").alias("loc_coarse"))
            for fn, kw in [(E.median_index, {"stat": "median"}), (E.stratified_median, {}), (E.pooled_td, {"robust": True}),
                           (E.rolling_td, {"window": 13, "splice": "mean"}), (E.hedonic_imputation, {"kind": "fisher"}), (E.geks, {"window": 25}),
                           (E.repeat_sales, {"weighted": True})]:
                t0 = time.time()
                try:
                    rec(rows, timing, cls, "region", reg, fn(drr, GRID, **kw), t0)
                except Exception as ex:
                    print("FAILED", cls, reg, fn.__name__, ex, flush=True)
            if reg in GBM_REGIONS and cls == "house":
                t0 = time.time(); rec(rows, timing, cls, "region", reg, E.gbm_residual_index(drr, GRID), t0)
            t0 = time.time()
            path, var, meth, gain = qhpi_cell(drr, qres.filter(pl.col("region_code") == reg), sK.log_index, GRID, qsig)
            rec(rows, timing, cls, "region", reg, E.IndexSeries(GRID, path, se=np.sqrt(var), n=E._counts(dr, GRID), method=f"K_qhpi_{meth}"), t0)
            if ("region", reg) in r.cells:
                p, v = cell_log_index(r, "region", reg, "smoothed")
                rec(rows, timing, cls, "region", reg, E.IndexSeries(GRID, p, se=np.sqrt(v), n=r.cells[("region", reg)].n, method="S_sharp"), time.time())
                p, v = cell_log_index(r, "region", reg, "filtered")
                rec(rows, timing, cls, "region", reg, E.IndexSeries(GRID, p, se=np.sqrt(v), n=r.cells[("region", reg)].n, method="S_sharp_filtered"), time.time())
        pl.concat(rows).write_parquet(OUT / "series.parquet")
        pl.DataFrame(timing).write_csv(OUT / "timing.csv")
    print("done")


if __name__ == "__main__":
    main()