#!/usr/bin/env python3
# =============================================================================
# QWHPI — Quebec Weekly Housing Price Index
# Author  : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# File    : engine/scripts/09_monthly.py
# Purpose : v2.1 — monthly hybrid indexes: RTD province anchor, DIRECT robust
#           local time-dummy for liquid cells, hierarchical shrinkage for
#           thin cells only.
# =============================================================================
"""Monthly index build (v2.1 headline).

Architecture (arbitrated against repeat sales, local time dummies and
stratified matched-cell medians — see methodology_v2_research.md):

- Province × type ....... rolling-time-dummy, Huber IRLS, mean splice.
- Liquid cell (median ≥ LIQUID_MIN tx/month) ... DIRECT local robust
  time-dummy (own β, own FSA FE), then a light state-space smoother
  (obs var σ²/nₘ) for the real-time variant. Deviation-from-pooled-surface
  methods compress strong local divergences (Québec City condo: +50–55%
  by every composition-free arbiter vs +42% by deviations), so liquid
  cells are estimated on their own data.
- Thin cell ............. deviation from its parent's published path via the
  heteroskedastic local-level Kalman model (unchanged machinery).
"""

from __future__ import annotations

import sys
import time
import unicodedata
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import numpy as np
import polars as pl

from qwhpi.clean import CLEAN_PARQUET, research_sample
from qwhpi.config import CITY_TARGETS, INTERIM_DIR, PROPERTY_TYPES, TABLES_DIR, ensure_dirs
from qwhpi.features import build_features
from qwhpi.hierarchy import CellDeviation, base_constant, cell_deviation, compose_cell_table
from qwhpi.rtd import fullpool_residualize, local_time_dummy, run_rtd
from qwhpi.state_space import fit_local_level

MONTHLY_PARQUET = INTERIM_DIR / "monthly_indexes.parquet"
FP_RESID = INTERIM_DIR / "rtd_residuals.parquet"
RTD_PATHS = INTERIM_DIR / "rtd_paths.parquet"

LIQUID_MIN = 40  # median tx/month for direct local estimation


def slug(name: str) -> str:
    s = unicodedata.normalize("NFKD", name).encode("ascii", "ignore").decode()
    return s.lower().replace(" ", "-").replace("'", "")


def city_slug(name: str) -> str:
    return "quebec-city" if name == "Québec" else slug(name)


def compose_direct(grid, raw_path, n, sigma2, gid, gname, glevel, ptype):
    """Published table for a directly-estimated liquid cell."""
    fit = fit_local_level(raw_path, np.maximum(n, 1.0), sigma2)
    base = base_constant(grid, fit.filtered)
    log_raw = raw_path - base
    log_filt = fit.filtered - base
    log_smooth = fit.smoothed - base
    se = np.sqrt(np.maximum(fit.filtered_var, 0.0))
    return pl.DataFrame({
        "week": grid,
        "geography_id": gid, "geography_name": gname, "geography_level": glevel,
        "property_type": ptype,
        "index": np.exp(log_raw) * 100.0,
        "index_smoothed": np.exp(log_filt) * 100.0,
        "index_research": np.exp(log_smooth) * 100.0,
        "log_index": log_filt,
        "se_log": se,
        "lower_95": np.exp(log_filt - 1.96 * se) * 100.0,
        "upper_95": np.exp(log_filt + 1.96 * se) * 100.0,
        "transactions": n,
        "shrinkage_weight": 1.0 - fit.gain,  # temporal smoothing weight
        "research_var": np.maximum(fit.smoothed_var, 0.0),
    }), fit


def combine_all(tables, weights, gid, gname, glevel):
    """Fixed-share composition of the three type indexes into 'all'."""
    wsum = sum(weights.values())
    w = {k: v / wsum for k, v in weights.items()}
    base = tables[next(iter(tables))].sort("week")
    grid = base["week"].to_list()

    def wavg(col, log=True):
        acc = np.zeros(len(grid))
        for k, tbl in tables.items():
            vals = tbl.sort("week")[col].to_numpy()
            acc += w[k] * (np.log(vals / 100.0) if log else vals)
        return acc

    log_raw, log_filt, log_res = wavg("index"), wavg("index_smoothed"), wavg("index_research")
    var = sum((w[k] ** 2) * (tbl.sort("week")["se_log"].to_numpy() ** 2)
              for k, tbl in tables.items())
    se = np.sqrt(var)
    n = sum(np.nan_to_num(tbl.sort("week")["transactions"].to_numpy())
            for tbl in tables.values())
    return pl.DataFrame({
        "week": grid,
        "geography_id": gid, "geography_name": gname, "geography_level": glevel,
        "property_type": "all",
        "index": np.exp(log_raw) * 100.0,
        "index_smoothed": np.exp(log_filt) * 100.0,
        "index_research": np.exp(log_res) * 100.0,
        "log_index": log_filt,
        "se_log": se,
        "lower_95": np.exp(log_filt - 1.96 * se) * 100.0,
        "upper_95": np.exp(log_filt + 1.96 * se) * 100.0,
        "transactions": n,
        "shrinkage_weight": wavg("shrinkage_weight", log=False),
        "research_var": var,
    })


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

    print("Stage 1: province RTD (13-month windows, Huber IRLS) ...")
    t0 = time.time()
    rtd = run_rtd(feat, verbose=False)
    print(f"  done in {time.time() - t0:.0f}s")
    rtd.drift.write_csv(TABLES_DIR / "rtd_coefficient_drift.csv")

    print("Stage 1b: full-pool robust residualization (thin-cell hierarchy) ...")
    fp_resid = fullpool_residualize(feat)
    fp_resid.write_parquet(FP_RESID)
    sigma2 = {t: float(fp_resid.filter(pl.col("propertyType") == t)["resid"].var())
              for t in PROPERTY_TYPES}

    grid = rtd.months
    T = len(grid)
    pl.concat([
        pl.DataFrame({"month": grid, "property_type": t,
                      "mean_splice": rtd.paths[t],
                      "movement_splice": rtd.paths_movement[t],
                      "window_splice": rtd.paths_window[t]})
        for t in PROPERTY_TYPES
    ]).write_parquet(RTD_PATHS)

    resid = fp_resid.rename({"month": "week_str"})
    region_names = dict(feat.select("region_code", "region").unique().iter_rows())
    frames: list[pl.DataFrame] = []
    estimation_log: list[dict] = []

    # Parent references for thin cells:
    region_resid_dev: dict[tuple[str, str], CellDeviation] = {}
    region_anchor: dict[tuple[str, str], np.ndarray] = {}  # published filtered log path

    for ptype in PROPERTY_TYPES:
        path = rtd.paths[ptype]
        r_t = resid.filter(pl.col("propertyType") == ptype)
        f_t = feat.filter(pl.col("propertyType") == ptype)
        counts = {r["week_str"]: r["n"] for r in
                  r_t.group_by("week_str").agg(pl.len().alias("n")).iter_rows(named=True)}
        n_prov = np.array([counts.get(m, 0) for m in grid], dtype=float)
        path_var = np.where(n_prov > 0, sigma2[ptype] / np.maximum(n_prov, 1), np.nan)
        path_var = np.nan_to_num(path_var, nan=float(np.nanmax(path_var)))

        prov = compose_cell_table(
            grid=grid, province_path=path, deviations=[],
            geography_id="quebec", geography_name="Québec (province)",
            geography_level="province", property_type=ptype, path_var=path_var,
        ).with_columns(pl.Series("transactions", n_prov),
                       pl.Series("shrinkage_weight", np.zeros(T)))
        frames.append(prov)
        estimation_log.append({"cell": f"quebec/{ptype}", "method": "rtd_mean_splice"})

        def monthly_counts(sub: pl.DataFrame) -> np.ndarray:
            c = dict(sub.group_by("month").agg(pl.len().alias("n")).iter_rows())
            return np.array([c.get(m, 0) for m in grid], dtype=float)

        # ---------------- regions ---------------- #
        for rcode in sorted(region_names):
            cell_r = r_t.filter(pl.col("region_code") == rcode)
            cell_f = f_t.filter(pl.col("region_code") == rcode)
            n_m = monthly_counts(cell_f)
            med = float(np.median(n_m))
            rdev = cell_deviation(cell_r.rename({"resid": "dev"})
                                  .select("week_str", "dev"), grid, sigma2[ptype])
            region_resid_dev[(rcode, ptype)] = rdev
            gid = f"region-{rcode}"

            direct = local_time_dummy(cell_f, grid) if med >= LIQUID_MIN else None
            if direct is not None:
                raw_path, n_cell, s2 = direct
                tbl, fit = compose_direct(grid, raw_path, n_cell, s2,
                                          gid, region_names[rcode], "region", ptype)
                frames.append(tbl)
                region_anchor[(rcode, ptype)] = raw_path * 0 + fit.filtered  # filtered log path
                estimation_log.append({"cell": f"{gid}/{ptype}", "method": "direct_local_td",
                                       "median_tx": med})
            else:
                tbl = compose_cell_table(
                    grid=grid, province_path=path, deviations=[rdev],
                    geography_id=gid, geography_name=region_names[rcode],
                    geography_level="region", property_type=ptype, path_var=path_var)
                frames.append(tbl)
                region_anchor[(rcode, ptype)] = path + rdev.fit.filtered
                estimation_log.append({"cell": f"{gid}/{ptype}", "method": "hierarchical",
                                       "median_tx": med})

        # ---------------- cities ---------------- #
        for city in CITY_TARGETS:
            cell_f = f_t.filter(pl.col("municipality") == city)
            if cell_f.height == 0:
                continue
            rcode = cell_f["region_code"][0]
            gid = city_slug(city)
            n_m = monthly_counts(cell_f)
            med = float(np.median(n_m))

            direct = local_time_dummy(cell_f, grid) if med >= LIQUID_MIN else None
            if direct is not None:
                raw_path, n_cell, s2 = direct
                tbl, _ = compose_direct(grid, raw_path, n_cell, s2,
                                        gid, city, "municipality", ptype)
                frames.append(tbl)
                estimation_log.append({"cell": f"{gid}/{ptype}", "method": "direct_local_td",
                                       "median_tx": med})
            else:
                # deviation of the city vs its REGION's residual path, anchored
                # on the region's published (filtered) log path
                rdev = region_resid_dev[(rcode, ptype)]
                rpath = {m: v for m, v in zip(grid, rdev.fit.smoothed)}
                cell_r = r_t.filter(pl.col("municipality") == city).with_columns(
                    (pl.col("resid") - pl.col("week_str").replace_strict(rpath, default=0.0))
                    .alias("dev"))
                mdev = cell_deviation(cell_r.select("week_str", "dev"), grid, sigma2[ptype])
                tbl = compose_cell_table(
                    grid=grid, province_path=region_anchor[(rcode, ptype)],
                    deviations=[mdev],
                    geography_id=gid, geography_name=city,
                    geography_level="municipality", property_type=ptype,
                    path_var=path_var)
                frames.append(tbl)
                estimation_log.append({"cell": f"{gid}/{ptype}", "method": "hierarchical",
                                       "median_tx": med})

    out = pl.concat(frames)

    all_frames = []
    for gid, gname, glevel in out.select(
            "geography_id", "geography_name", "geography_level").unique().iter_rows():
        tables = {t: out.filter((pl.col("geography_id") == gid)
                                & (pl.col("property_type") == t))
                  for t in PROPERTY_TYPES}
        tables = {t: tbl for t, tbl in tables.items() if tbl.height}
        weights = {t: max(float(np.nansum(tbl["transactions"].to_numpy())), 1.0)
                   for t, tbl in tables.items()}
        all_frames.append(combine_all(tables, weights, gid, gname, glevel))

    out = pl.concat([out, *all_frames]).rename({"week": "period"}).sort(
        ["geography_level", "geography_id", "property_type", "period"])
    out.write_parquet(MONTHLY_PARQUET)

    log_df = pl.DataFrame(estimation_log)
    log_df.write_csv(TABLES_DIR / "estimation_methods.csv")
    n_direct = log_df.filter(pl.col("method") == "direct_local_td").height
    n_hier = log_df.filter(pl.col("method") == "hierarchical").height
    print(f"\nCells: {n_direct} direct, {n_hier} hierarchical, 3 RTD province")
    print(f"Saved {out.height:,} observations to {MONTHLY_PARQUET}")


if __name__ == "__main__":
    main()