#!/usr/bin/env python3
# =============================================================================
# QWHPI — Quebec Weekly Housing Price Index
# Author  : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# File    : engine/scripts/04_baseline.py
# Purpose : Step 4 — pooled hedonic time-dummy baseline for Quebec + the four
#           major cities × All/Unifamilial/Condo/Plex (20 cells).
# =============================================================================
"""Baseline index build (Execution Order step 4).

For every cell: pooled Model A time-dummy regression (2021→present),
week coefficients rebased to 2021 avg = 100, raw weekly median attached for
comparison. Results land in ``data/interim/baseline_indexes.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 qwhpi.clean import CLEAN_PARQUET, research_sample
from qwhpi.config import INTERIM_DIR, TABLES_DIR, ensure_dirs
from qwhpi.features import build_features
from qwhpi.hedonic import estimate_time_dummy

BASELINE_PARQUET = INTERIM_DIR / "baseline_indexes.parquet"

GEOGRAPHIES = [
    ("quebec", "Québec (province)", "province", None),
    ("montreal", "Montréal", "municipality", "Montréal"),
    ("quebec-city", "Québec", "municipality", "Québec"),
    ("laval", "Laval", "municipality", "Laval"),
    ("gatineau", "Gatineau", "municipality", "Gatineau"),
]
TYPE_SETS = ["all", "unifamilial", "condo", "plex"]


def main() -> None:
    ensure_dirs()
    clean = pl.read_parquet(CLEAN_PARQUET)
    feat = build_features(research_sample(clean))

    frames: list[pl.DataFrame] = []
    summary_rows: list[dict] = []

    for geo_id, geo_name, level, mun_filter in GEOGRAPHIES:
        base = feat if mun_filter is None else feat.filter(pl.col("municipality") == mun_filter)
        for ptype in TYPE_SETS:
            cell = base if ptype == "all" else base.filter(pl.col("propertyType") == ptype)
            t0 = time.time()
            res = estimate_time_dummy(cell, absorb="loc_fine",
                                      include_property_type=(ptype == "all"))
            if res is None:
                print(f"[skip] {geo_id} x {ptype}: n={cell.height} too thin")
                summary_rows.append({"geography": geo_id, "property_type": ptype,
                                     "n_obs": cell.height, "estimated": False,
                                     "r2": None, "median_weekly_tx": None})
                continue

            # Raw weekly median + volume for comparison/validation.
            weekly_raw = (
                cell.group_by("week_str")
                .agg(pl.col("amount").median().alias("raw_median"),
                     pl.len().alias("transactions"))
                .rename({"week_str": "week"})
            )
            tbl = (
                res.table.join(weekly_raw, on="week", how="left")
                .with_columns(
                    pl.lit(geo_id).alias("geography_id"),
                    pl.lit(geo_name).alias("geography_name"),
                    pl.lit(level).alias("geography_level"),
                    pl.lit(ptype).alias("property_type"),
                )
            )
            frames.append(tbl)
            med_tx = weekly_raw["transactions"].median()
            summary_rows.append({"geography": geo_id, "property_type": ptype,
                                 "n_obs": res.n_obs, "estimated": True,
                                 "r2": round(res.r2, 4),
                                 "median_weekly_tx": med_tx})
            print(f"[ok] {geo_id} x {ptype}: n={res.n_obs:,} r2={res.r2:.3f} "
                  f"({time.time() - t0:.1f}s)")

    out = pl.concat(frames)
    out.write_parquet(BASELINE_PARQUET)
    pl.DataFrame(summary_rows).write_csv(TABLES_DIR / "baseline_model_summary.csv")
    print(f"\nSaved {out.height:,} index observations to {BASELINE_PARQUET}")


if __name__ == "__main__":
    main()