#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 2 — HMTS benchmark (G) at province × class and region × class, windows 4 and 13 months.
Output: results/benchmarks/series_hmts.parquet (same long schema as series.parquet)."""
from __future__ import annotations

import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
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

OUT = C.RESULTS / "benchmarks"
GRID = E.month_grid("2021-01", "2026-10")


def main() -> None:
    f = build_features(pl.read_parquet(C.ANALYTIC_TX))
    rows = []
    for cls in C.PROPERTY_TYPES:
        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"))
        for w in (4, 13):
            t0 = time.time(); s = E.hmts(prov, GRID, window=w)
            rows.append(s.to_frame().with_columns(pl.lit(cls).alias("class"), pl.lit("province").alias("level"), pl.lit("QC").alias("geo_id"), pl.lit(None, dtype=pl.Float64).alias("se")))
            print(cls, "province", s.method, f"{time.time() - t0:.0f}s", flush=True)
        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"))
            t0 = time.time(); s = E.hmts(drr, GRID, window=13)
            rows.append(s.to_frame().with_columns(pl.lit(cls).alias("class"), pl.lit("region").alias("level"), pl.lit(reg).alias("geo_id"), pl.lit(None, dtype=pl.Float64).alias("se")))
            print(cls, reg, s.method, f"{time.time() - t0:.0f}s", flush=True)
        pl.concat(rows).select(["class", "level", "geo_id", "method", "month", "log_index", "se", "n"]).write_parquet(OUT / "series_hmts.parquet")


if __name__ == "__main__":
    main()