# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Arm A: hedonic ridge regression on log price with municipality x quarter and 1-km grid fixed effects.
Arm B: hierarchical hedonic (empirical-Bayes shrinkage of region > municipality > grid residual means)
with a repeat-sales adjustment when the parcel has a prior sale."""
from __future__ import annotations

import numpy as np
import polars as pl
import scipy.sparse as sp
from sklearn.linear_model import Ridge
from sklearn.preprocessing import OneHotEncoder, StandardScaler

from .common import Prediction, QUANTILES, numeric_features, residual_quantile_intervals


def _quarter(df: pl.DataFrame) -> pl.Series:
    return (df["date"].dt.year().cast(pl.Utf8) + "Q" + ((df["date"].dt.month() - 1) // 3 + 1).cast(pl.Utf8))


class HedonicRidge:
    def __init__(self, alpha: float = 1.0, fe: tuple[str, ...] = ("mun_code", "grid1k", "property_class", "subtype", "mun_class")):
        # Out-of-time prediction: calendar-quarter fixed effects (and polynomial trends) cannot be extrapolated
        # safely to 2025-26 (a cubic trend fitted on 2021-24 produced a +0.5 log bias on the 2026 holdout). Time
        # therefore enters only through month-of-year effects and the as-of market-level features of the matrix
        # (12-month municipality and 3-km neighbourhood medians), which are legitimately observable at the sale date.
        self.alpha = alpha
        self.fe = fe
        self.num_cols: list[str] = []
        self.enc = OneHotEncoder(handle_unknown="ignore", min_frequency=5)
        self.scaler = StandardScaler()
        self.model = Ridge(alpha=alpha, solver="sparse_cg", random_state=0)
        self.resid_valid: np.ndarray | None = None

    def _design(self, df: pl.DataFrame, fit: bool) -> sp.csr_matrix:
        d = df.with_columns(
            (pl.col("mun_code").cast(pl.Utf8) + "_" + pl.col("property_class").cast(pl.Utf8)).alias("mun_class"),
            ((pl.col("date").dt.year() - 2021) * 12 + pl.col("date").dt.month() - 1).cast(pl.Float64).alias("_t"),
            pl.col("date").dt.month().cast(pl.Utf8).alias("_moy"),
        )
        cat = d.select([c for c in self.fe if c in d.columns] + ["_moy"]).fill_null("NA").to_pandas().astype(str)
        if fit:
            self.num_cols = [c for c in numeric_features(df) if c not in ("lat", "lng", "x_m", "y_m", "month_index")]
        X = d.select(self.num_cols).fill_null(strategy="mean").fill_nan(0.0).to_numpy().astype(float)
        X = np.nan_to_num(X)
        self._mun_trend = sp.csr_matrix((len(X), 0))
        if fit:
            Xc = self.enc.fit_transform(cat)
            Xn = self.scaler.fit_transform(X)
        else:
            Xc = self.enc.transform(cat)
            Xn = self.scaler.transform(X)
        return sp.hstack([sp.csr_matrix(Xn), Xc, self._mun_trend]).tocsr()

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame | None = None) -> "HedonicRidge":
        X = self._design(train, fit=True)
        self.model.fit(X, train["log_price"].to_numpy())
        if valid is not None:
            self.resid_valid = valid["log_price"].to_numpy() - self.model.predict(self._design(valid, fit=False))
        return self

    def predict_log(self, df: pl.DataFrame) -> np.ndarray:
        return self.model.predict(self._design(df, fit=False))

    def predict(self, df: pl.DataFrame) -> Prediction:
        p = self.predict_log(df)
        q = residual_quantile_intervals(p, self.resid_valid) if self.resid_valid is not None else None
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(p), quantiles=q)


class HierarchicalHedonic:
    """Ridge hedonic without geographic fixed effects + empirical-Bayes shrinkage of residual means at region,
    municipality and 1-km grid levels (estimated on train), + repeat-sales blend when a prior sale exists.

    The blend weight w on the repeat-sales forecast (prior log price + municipality index change) is chosen on
    validation by least squares, never on test."""

    def __init__(self, alpha: float = 1.0, k_shrink: float = 20.0):
        self.base = HedonicRidge(alpha=alpha, fe=("property_class", "subtype"))
        self.k = k_shrink
        self.levels = ("region_name", "mun_code", "grid1k")
        self.effects: dict[str, dict[str, float]] = {}
        self.w_repeat = 0.0
        self.resid_valid: np.ndarray | None = None

    def _shrunk_effects(self, df: pl.DataFrame, resid: np.ndarray) -> np.ndarray:
        r = resid.copy()
        total = np.zeros_like(r)
        for lvl in self.levels:
            key = df[lvl].cast(pl.Utf8).fill_null("NA").to_numpy()
            tab = pl.DataFrame({"k": key, "r": r}).group_by("k").agg(pl.col("r").mean().alias("m"), pl.len().alias("n"))
            eff = {k: m * n / (n + self.k) for k, m, n in zip(tab["k"], tab["m"], tab["n"])}
            self.effects[lvl] = eff
            e = np.array([eff.get(k, 0.0) for k in key])
            r = r - e
            total += e
        return total

    def _apply_effects(self, df: pl.DataFrame) -> np.ndarray:
        total = np.zeros(df.height)
        for lvl in self.levels:
            key = df[lvl].cast(pl.Utf8).fill_null("NA").to_numpy()
            total += np.array([self.effects[lvl].get(k, 0.0) for k in key])
        return total

    def _repeat_forecast(self, df: pl.DataFrame) -> np.ndarray:
        if "prior_sale_log_price" in df.columns and "prior_sale_index_change" in df.columns:
            return df.select((pl.col("prior_sale_log_price") + pl.col("prior_sale_index_change").fill_null(0.0))).to_numpy().ravel()
        if "prior_sale_log_price" in df.columns:
            return df["prior_sale_log_price"].to_numpy()
        return np.full(df.height, np.nan)

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame | None = None) -> "HierarchicalHedonic":
        self.base.fit(train)
        resid = train["log_price"].to_numpy() - self.base.predict_log(train)
        self._shrunk_effects(train, resid)
        if valid is not None:
            h = self.base.predict_log(valid) + self._apply_effects(valid)
            rs = self._repeat_forecast(valid)
            y = valid["log_price"].to_numpy()
            m = ~np.isnan(rs)
            if m.sum() > 50:
                d = rs[m] - h[m]
                self.w_repeat = float(np.clip(np.sum(d * (y[m] - h[m])) / max(np.sum(d * d), 1e-9), 0.0, 1.0))
            self.resid_valid = y - self._combine(h, rs)
        return self

    def _combine(self, h: np.ndarray, rs: np.ndarray) -> np.ndarray:
        out = h.copy()
        m = ~np.isnan(rs)
        out[m] = (1 - self.w_repeat) * h[m] + self.w_repeat * rs[m]
        return out

    def predict_log(self, df: pl.DataFrame) -> np.ndarray:
        return self._combine(self.base.predict_log(df) + self._apply_effects(df), self._repeat_forecast(df))

    def predict(self, df: pl.DataFrame) -> Prediction:
        p = self.predict_log(df)
        q = residual_quantile_intervals(p, self.resid_valid) if self.resid_valid is not None else None
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(p), quantiles=q)