# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Benchmark K — re-implementation of the QWHPI v2.1 monthly architecture (PhD/recherche/QWHPI, 2026-08):
province × class = 13-month rolling time dummy, Huber IRLS, mean splice; liquid cells (median ≥ 40 tx/month) = direct
local time dummy lightly smoothed by a local-level filter; thin cells = local-level Kalman deviation from the parent's
published path with per-cell maximum-likelihood τ² (no pooling across cells)."""
from __future__ import annotations

import numpy as np
import polars as pl

from sharp_hpi.estimators import IndexSeries, pooled_td, rolling_td
from sharp_hpi.hedonic import RHS_STRUCT, assign_loc_fe, fit_fe
from sharp_hpi.statespace import fit_cell

LIQUID_MIN = 40


def qhpi_province(df: pl.DataFrame, grid: list[str]) -> IndexSeries:
    d = assign_loc_fe(df, min_n=15)
    return rolling_td(d, grid, window=13, splice="mean", robust=True)


def qhpi_fullpool_residuals(df: pl.DataFrame, grid: list[str]) -> tuple[pl.DataFrame, float]:
    d = assign_loc_fe(df, min_n=15)
    rhs = RHS_STRUCT.replace(" + C(units_c)", "") if d["property_class"][0] != "plex" else RHS_STRUCT
    if d["property_class"][0] != "house":
        rhs = rhs.replace(" + C(subtype_c)", "")
    res, m = fit_fe(d, rhs + " + C(month)", "loc_fe", robust=True, n_irls=2)
    return d.select(["tx_id", "month", "region_code", "mrc_code", "mun_code"]).with_columns(pl.Series("e", res.resid)), res.sigma2_robust


def qhpi_cell(df_cell: pl.DataFrame, resid_cell: pl.DataFrame, parent_path: np.ndarray, grid: list[str], sigma2: float) -> tuple[np.ndarray, np.ndarray, str, np.ndarray]:
    """Returns (log path, variance, method, gain) for a sub-provincial cell."""
    cnt = resid_cell.group_by("month").agg(pl.len().alias("n"))
    nmap = dict(zip(cnt["month"].to_list(), cnt["n"].to_list()))
    n = np.array([nmap.get(m, 0) for m in grid], float)
    T = len(grid)
    if np.median(n[n > 0]) >= LIQUID_MIN if (n > 0).any() else False:
        d = assign_loc_fe(df_cell, min_n=15)
        s = pooled_td(d, grid, robust=True)
        y = s.log_index
        ok = ~np.isnan(y)
        # light smoothing: local level with obs var sigma2/n on the direct TD path
        f = fit_cell(np.where(ok, y, np.nan), np.where(ok, n, 0), s.meta.get("sigma_robust", np.sqrt(sigma2)) ** 2, model="ll", robust=False)
        return f.smoothed, f.smoothed_var, "direct_local_td", f.gain
    agg = resid_cell.group_by("month").agg(pl.col("e").mean().alias("yb"))
    ymap = dict(zip(agg["month"].to_list(), agg["yb"].to_list()))
    y = np.array([ymap.get(m, np.nan) for m in grid])
    f = fit_cell(y, n, sigma2, model="ll", robust=False)
    return parent_path + f.smoothed, f.smoothed_var, "hierarchical", f.gain