# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Hedonic engine: feature construction, robust (Huber-IRLS) fixed-effects regressions.

Headline specification (structural + location only; no market variables):
    log P = s(log floor area) + s(log land area) + age bins + storeys + genre + subtype [+ n_units for plex]
            + location FE (fine: neighbourhood unit / 1-km cell)  + time effects
Missing characteristics: conditional-median imputation within (class × municipality) → (class × region) → class,
with time-stable missingness indicators (floor-area and land-area missingness are checked for drift in
results/profile/completeness_by_month.csv).
"""
from __future__ import annotations

from dataclasses import dataclass

import warnings

import numpy as np
import polars as pl
import pyfixest as pf

warnings.filterwarnings("ignore", message=".*singleton fixed effect.*")
warnings.filterwarnings("ignore", message=".*multicollinearity.*")
warnings.filterwarnings("ignore", category=RuntimeWarning, module="sharp_hpi")

AGE_EDGES = [1, 5, 10, 20, 40, 60, 80, 120]
AGE_LABELS = ["a00_0-1", "a01_2-5", "a02_6-10", "a03_11-20", "a04_21-40", "a05_41-60", "a06_61-80", "a07_81-120", "a08_120+"]
HUBER_K = 1.345
N_IRLS = 3


def _impute(df: pl.DataFrame, col: str, filled: str, cls: str = "property_class") -> pl.DataFrame:
    m1 = df.group_by([cls, "mun_code"]).agg(pl.col(col).median().alias("_m1"))
    m2 = df.group_by([cls, "region_code"]).agg(pl.col(col).median().alias("_m2"))
    m3 = df.group_by(cls).agg(pl.col(col).median().alias("_m3"))
    return (df.join(m1, on=[cls, "mun_code"], how="left").join(m2, on=[cls, "region_code"], how="left").join(m3, on=cls, how="left")
            .with_columns(pl.coalesce(pl.col(col), pl.col("_m1"), pl.col("_m2"), pl.col("_m3")).alias(filled)).drop("_m1", "_m2", "_m3"))


def build_features(df: pl.DataFrame) -> pl.DataFrame:
    """Attach modelling features to the analytic layer (idempotent)."""
    out = df.with_columns(
        pl.when(pl.col("floor_area_m2") > 10).then(pl.col("floor_area_m2")).otherwise(None).alias("_fa"),
        pl.when(pl.col("land_area_m2") > 10).then(pl.col("land_area_m2")).otherwise(None).alias("_la"),
        pl.when(pl.col("age").is_null() | (pl.col("age") > 200)).then(None).otherwise(pl.col("age")).alias("_age"),
    ).with_columns(
        pl.col("_fa").is_null().alias("fa_missing"), pl.col("_la").is_null().alias("la_missing"), pl.col("_age").is_null().alias("age_missing"),
    )
    out = _impute(out, "_fa", "fa_f")
    out = _impute(out, "_la", "la_f")
    out = _impute(out, "_age", "age_f")
    age = pl.col("age_f")
    age_bin = pl.when(age <= 1).then(pl.lit(AGE_LABELS[0]))
    lo = 1
    for edge, lab in zip(AGE_EDGES[1:], AGE_LABELS[1:-1]):
        age_bin = age_bin.when((age > lo) & (age <= edge)).then(pl.lit(lab))
        lo = edge
    age_bin = age_bin.otherwise(pl.lit(AGE_LABELS[-1]))
    out = out.with_columns(
        pl.col("fa_f").log().alias("log_fa"), pl.col("la_f").log().alias("log_la"), age_bin.alias("age_bin"),
        pl.col("storeys").fill_null(1.0).clip(1, 4).round(0).cast(pl.Int32).cast(pl.Utf8).alias("storeys_c"),
        pl.col("subtype").fill_null("na").alias("subtype_c"),
        pl.when(pl.col("property_class") == "plex").then(pl.col("n_units").clip(2, 5)).otherwise(1).cast(pl.Utf8).alias("units_c"),
        pl.col("genre").alias("genre_c"),
        pl.concat_str([pl.col("property_class"), pl.lit(":"), pl.col("month")]).alias("class_month"),
    ).drop("_fa", "_la", "_age")
    # log floor area: linear + squared + cubic around the class median (smooth nonlinearity); land area linear + squared
    med = out.group_by("property_class").agg(pl.col("log_fa").median().alias("_mfa"), pl.col("log_la").median().alias("_mla"))
    out = out.join(med, on="property_class", how="left").with_columns(
        (pl.col("log_fa") - pl.col("_mfa")).alias("lfa"), (pl.col("log_la") - pl.col("_mla")).alias("lla"),
    ).with_columns((pl.col("lfa") ** 2).alias("lfa2"), (pl.col("lfa") ** 3).alias("lfa3"), (pl.col("lla") ** 2).alias("lla2")).drop("_mfa", "_mla")
    return out


def assign_loc_fe(df: pl.DataFrame, min_n: int = 15, levels: tuple[str, ...] = ("grid500", "grid1k", "grid2k", "grid5k", "mun_code", "mrc_code"),
                  group: tuple[str, ...] = ("property_class",)) -> pl.DataFrame:
    """Finest nested spatial cell with at least `min_n` rows (within `group`), coarsening upward.

    The partition depends only on sale locations (not prices) and is fixed for the estimation sample it is built on;
    real-time vintages rebuild it from the data available at the cutoff.
    """
    g = list(group)
    key = pl.lit(None, dtype=pl.Utf8)
    chosen = pl.lit(None, dtype=pl.Utf8)
    out = df
    for lvl in levels:
        n = pl.len().over(g + [lvl])
        chosen = pl.when(chosen.is_null() & (n >= min_n)).then(pl.concat_str([pl.lit(lvl + ":"), pl.col(lvl)])).otherwise(chosen)
    out = out.with_columns(chosen.alias("loc_fe"))
    out = out.with_columns(pl.when(pl.col("loc_fe").is_null()).then(pl.concat_str([pl.lit("mrc_code:"), pl.col("mrc_code")])).otherwise(pl.col("loc_fe")).alias("loc_fe"))
    # a cell that received rows only through fall-back might still be < min_n at a coarse level; that is acceptable (no singletons by
    # construction only if the group has >= min_n rows at the top level) — guard below.
    return out.with_columns(pl.col("loc_fe").str.slice(0, 8).alias("loc_fe_level"))


# Right-hand sides -----------------------------------------------------------
RHS_STRUCT = "lfa + lfa2 + lfa3 + fa_missing + lla + lla2 + la_missing + C(age_bin) + age_missing + C(storeys_c) + C(genre_c) + C(subtype_c) + C(units_c)"
RHS_STRUCT_BY_CLASS = "lfa + lfa2 + lfa3 + fa_missing + lla + lla2 + la_missing + C(age_bin) + age_missing + C(storeys_c) + C(genre_c) + C(subtype_c) + C(units_c)"


@dataclass
class FitResult:
    coef: dict
    resid: np.ndarray
    fitted: np.ndarray
    weights: np.ndarray
    sigma2_robust: float
    n: int
    r2_within: float | None


def huber_weights(resid: np.ndarray, k: float = HUBER_K) -> tuple[np.ndarray, float]:
    s = 1.4826 * np.median(np.abs(resid - np.median(resid)))
    s = max(s, 1e-6)
    u = np.abs(resid) / s
    w = np.where(u <= k, 1.0, k / u)
    return w, s


def fit_fe(df: pl.DataFrame, rhs: str, fe: str, y: str = "log_price", robust: bool = True, n_irls: int = N_IRLS,
           base_weights: np.ndarray | None = None) -> tuple[FitResult, pf.estimation.Feols]:
    """Huber-IRLS fixed-effects regression (pyfixest). Returns residuals aligned with df rows."""
    n_all = df.height
    w_all = np.ones(n_all) if base_weights is None else np.asarray(base_weights, float)
    # iterative singleton removal on every FE factor so residuals stay aligned (dropped rows get NaN)
    keep = np.ones(n_all, bool)
    fe_cols = [c.strip() for c in fe.split("+")] if fe else []
    if fe_cols:
        sub = df.select(fe_cols).with_row_index("_i")
        while True:
            cur = sub.filter(pl.Series(keep)[sub["_i"]])
            drop = np.zeros(n_all, bool)
            for c in fe_cols:
                cnt = cur.group_by(c).agg(pl.len().alias("_n"))
                ones = cur.join(cnt.filter(pl.col("_n") == 1), on=c, how="semi")["_i"].to_numpy()
                drop[ones] = True
            if not drop.any():
                break
            keep &= ~drop
    pdf = df.filter(pl.Series(keep)).to_pandas()
    w = w_all[keep]
    fml = f"{y} ~ {rhs} | {fe}" if fe else f"{y} ~ {rhs}"
    pdf["_w"] = w
    m = pf.feols(fml, data=pdf, weights="_w", vcov="iid")
    resid = np.asarray(m.resid(), float)
    if len(resid) != len(pdf):
        raise ValueError(f"rows dropped by the FE regression ({len(pdf) - len(resid)}) after singleton removal")
    if robust:
        for _ in range(n_irls):
            hw, _s = huber_weights(resid)
            pdf["_w"] = w * hw
            m = pf.feols(fml, data=pdf, weights="_w", vcov="iid")
            resid = np.asarray(m.resid(), float)
    hw, s = huber_weights(resid)
    resid_all = np.full(n_all, np.nan); resid_all[keep] = resid
    fitted_all = np.full(n_all, np.nan); fitted_all[keep] = np.asarray(pdf[y].values, float) - resid
    w_out = np.zeros(n_all); w_out[keep] = pdf["_w"].values
    r2w = float(m._r2_within) if hasattr(m, "_r2_within") and m._r2_within is not None else None
    return FitResult(coef=dict(m.coef()), resid=resid_all, fitted=fitted_all, weights=w_out, sigma2_robust=float(s ** 2), n=int(keep.sum()), r2_within=r2w), m