# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""SHARP-HPI: Spatially Hierarchical, Adaptive and Robust housing price index.

Decomposition (per class k):
    log P_i = f(X_i; θ_t) + α_loc(i) + μ_{k,t} + δ^R_{r(i),t} + δ^M_{m(i),t} + δ^U_{u(i),t} + ε_i

Stage 1 — robust pooled hedonic (Huber-IRLS) with fine-location fixed effects α_loc (nested metric grid),
          province × month effects μ_t, and smoothly time-varying quality prices (interactions of the
          structural terms with a centred linear time index). Residuals e_i are quality- and
          location-adjusted log prices net of the provincial path.
Stage 2 — nested deviations. For each level ℓ ∈ {region, MRC, municipality} the monthly mean residual of a
          cell (net of its parent's smoothed/filtered deviation) is modelled by a heteroskedastic local-level
          (or local-linear-trend) state-space model with observation variance σ²_cell / n_t. Innovation
          variances τ² are calibrated by empirical Bayes across the cells of a level (grid EM on a normal prior for
          log τ²), so a thin cell's responsiveness is borrowed from its peers instead of being estimated from
          its own handful of sales (the adaptive pooling). Innovations are Huber down-weighted (robust).
Real-time mode uses the filtered path at every level; retrospective mode the smoothed path.

Reliability bundle per cell-month: actual count, effective sample size (σ²_cell / Var(state)), local data weight
(Kalman gain), posterior sd, covariate overlap, residual autocorrelation, one-step-ahead standardized error.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import polars as pl

from sharp_hpi.hedonic import RHS_STRUCT, assign_loc_fe, fit_fe
from sharp_hpi.statespace import LOG_TAU2_GRID, SSFit, eb_prior, fit_cell, profile_loglik

LEVELS = [("region", "region_code"), ("mrc", "mrc_code"), ("municipality", "mun_code")]
TV_TERMS = ["lfa", "lfa2", "lla"]  # quality prices allowed to drift linearly in time (ablation: none)


@dataclass
class Stage1:
    months: list[str]
    mu: np.ndarray            # provincial log path (month effects, first month = 0)
    mu_se: np.ndarray
    resid: pl.DataFrame       # tx_id, month, region_code, mrc_code, mun_code, e
    sigma2: float
    coef: dict
    n: int
    r2_within: float | None


@dataclass
class CellFit:
    level: str
    geo_id: str
    parent_id: str | None
    months: list[str]
    n: np.ndarray
    ybar: np.ndarray          # observed monthly mean deviation (NaN if none)
    fit: SSFit
    sigma2_cell: float
    overlap: np.ndarray       # share of the cell's sales inside the class common support
    rho1: float               # lag-1 autocorrelation of observed deviations (residual dependence)
    oos_z: np.ndarray         # one-step-ahead standardized innovation (NaN if none)


@dataclass
class SharpResult:
    cls: str
    stage1: Stage1
    cells: dict[tuple[str, str], CellFit] = field(default_factory=dict)
    priors: dict[str, tuple[float, float]] = field(default_factory=dict)
    config: dict = field(default_factory=dict)


def _rhs(df: pl.DataFrame, time_varying: bool) -> str:
    rhs = RHS_STRUCT
    cls = df["property_class"][0]
    if cls != "plex":
        rhs = rhs.replace(" + C(units_c)", "")
    if cls != "house":
        rhs = rhs.replace(" + C(subtype_c)", "")
    if time_varying:
        rhs += " + " + " + ".join(f"{t}:tt" for t in TV_TERMS)
    return rhs + " + C(month)"


def stage1(df: pl.DataFrame, grid: list[str], time_varying: bool = True, min_loc_n: int = 15, robust: bool = True, n_irls: int = 2) -> Stage1:
    T = len(grid)
    tt = {m: (i - (T - 1) / 2) / max(T - 1, 1) for i, m in enumerate(grid)}
    d = assign_loc_fe(df, min_n=min_loc_n).with_columns(pl.col("month").replace_strict(tt, default=0.0).alias("tt"))
    res, m = fit_fe(d, _rhs(d, time_varying), "loc_fe", robust=robust, n_irls=n_irls)
    se = dict(m.se())
    base = grid[0]
    mu = np.array([0.0 if g == base else res.coef.get(f"C(month)[T.{g}]", np.nan) for g in grid])
    mu_se = np.array([0.0 if g == base else se.get(f"C(month)[T.{g}]", np.nan) for g in grid])
    resid = d.select(["tx_id", "month", "region_code", "mrc_code", "mun_code", "lfa", "age_f"]).with_columns(pl.Series("e", res.resid))
    return Stage1(grid, mu, mu_se, resid, res.sigma2_robust, res.coef, res.n, res.r2_within)


def _cell_series(df: pl.DataFrame, key: str, grid: list[str], ycol: str = "e") -> dict[str, tuple[np.ndarray, np.ndarray, float, np.ndarray]]:
    """Per cell id: (ybar by month, n by month, robust sigma2, overlap by month)."""
    agg = df.group_by([key, "month"]).agg(pl.col(ycol).mean().alias("ybar"), pl.len().alias("n"), pl.col("_in_support").mean().alias("ov"))
    sig = df.group_by(key).agg(((pl.col(ycol) - pl.col(ycol).median()).abs().median() * 1.4826).pow(2).alias("s2"), pl.len().alias("N"))
    idx = {m: i for i, m in enumerate(grid)}
    out = {}
    s2 = dict(zip(sig[key].to_list(), zip(sig["s2"].to_list(), sig["N"].to_list())))
    for gid, sub in agg.group_by(key):
        gid = gid[0]
        y = np.full(len(grid), np.nan); n = np.zeros(len(grid)); ov = np.full(len(grid), np.nan)
        for mth, yb, nn, o in zip(sub["month"].to_list(), sub["ybar"].to_list(), sub["n"].to_list(), sub["ov"].to_list()):
            if mth in idx:
                y[idx[mth]] = yb; n[idx[mth]] = nn; ov[idx[mth]] = o
        out[gid] = (y, n, s2[gid], ov)
    return out


def _acf1(y: np.ndarray) -> float:
    v = y[~np.isnan(y)]
    if len(v) < 6:
        return np.nan
    v = v - v.mean()
    return float(np.sum(v[1:] * v[:-1]) / np.sum(v * v)) if np.sum(v * v) > 0 else np.nan


def stage2(s1: Stage1, df_geo: pl.DataFrame, grid: list[str], model: str = "ll", eb: bool = True, robust: bool = True,
           realtime_parent: bool = False, sigma2_shrink_n: int = 100, levels=LEVELS) -> tuple[dict, dict]:
    """Nested state-space deviations. Returns (cells, priors)."""
    e = s1.resid
    # common support: class-level 1–99 % band of (centred) log floor area and age
    q = e.select(pl.col("lfa").quantile(0.01).alias("l1"), pl.col("lfa").quantile(0.99).alias("l99"), pl.col("age_f").quantile(0.99).alias("a99"))
    e = e.with_columns((pl.col("lfa").is_between(q["l1"][0], q["l99"][0]) & (pl.col("age_f") <= q["a99"][0])).alias("_in_support"))
    e = e.with_columns(pl.lit(0.0).alias("parent_dev"))
    cells: dict[tuple[str, str], CellFit] = {}
    priors: dict[str, tuple[float, float]] = {}
    parent_key = None
    for lname, key in levels:
        e = e.with_columns((pl.col("e") - pl.col("parent_dev")).alias("e_net"))
        series = _cell_series(e, key, grid, "e_net")
        # cell sigma2 shrunk toward the class robust variance for small cells
        for gid, (y, n, (s2, N), ov) in series.items():
            w = N / (N + sigma2_shrink_n)
            series[gid] = (y, n, (w * s2 + (1 - w) * s1.sigma2, N), ov)
        prior = None
        if eb:
            profs = [profile_loglik(y, n, s2[0], model) for (y, n, s2, ov) in series.values() if (n > 0).sum() >= 6]
            prior = eb_prior(profs) if len(profs) >= 3 else None
            priors[lname] = prior if prior else (float(np.log(1e-5)), 1.0)
        dev_map = {}
        for gid, (y, n, (s2, N), ov) in series.items():
            fit = fit_cell(y, n, s2, model=model, prior=prior if eb else None, robust=robust)
            parent = e.filter(pl.col(key) == gid)[parent_key][0] if parent_key else None
            z = np.where(n > 0, (np.nan_to_num(y) - np.concatenate([[0.0], fit.filtered[:-1]])) / np.sqrt(np.where(n > 0, s2 / np.maximum(n, 1), 1) + fit.tau2 + np.concatenate([[0.25], fit.filtered_var[:-1]])), np.nan)
            cells[(lname, gid)] = CellFit(lname, gid, parent, grid, n, y, fit, s2, ov, _acf1(y), z)
            dev_map[gid] = fit.filtered if realtime_parent else fit.smoothed
        # pass this level's deviation to the children as their parent path
        idx = {m: i for i, m in enumerate(grid)}
        dev_tbl = pl.DataFrame({key: list(dev_map.keys()), "_dev": [list(v) for v in dev_map.values()]}).explode("_dev").with_columns(
            pl.Series("month", grid * len(dev_map)))
        e = e.drop("parent_dev").join(dev_tbl, on=[key, "month"], how="left").with_columns(pl.col("_dev").fill_null(0.0).alias("_dev"))
        e = e.with_columns((pl.col("parent_dev_prev") if "parent_dev_prev" in e.columns else pl.lit(0.0)).alias("_pp"))
        e = e.with_columns((pl.col("_dev") + pl.col("_pp")).alias("parent_dev")).drop("_dev", "_pp")
        e = e.with_columns(pl.col("parent_dev").alias("parent_dev_prev"))
        parent_key = key
    return cells, priors


def run_sharp(df: pl.DataFrame, grid: list[str], time_varying: bool = True, model: str = "ll", eb: bool = True, robust: bool = True,
              realtime_parent: bool = False, min_loc_n: int = 15, levels=LEVELS) -> SharpResult:
    cls = df["property_class"][0]
    s1 = stage1(df, grid, time_varying=time_varying, min_loc_n=min_loc_n, robust=robust)
    cells, priors = stage2(s1, df, grid, model=model, eb=eb, robust=robust, realtime_parent=realtime_parent, levels=levels)
    return SharpResult(cls, s1, cells, priors, {"time_varying": time_varying, "model": model, "eb": eb, "robust": robust, "realtime_parent": realtime_parent, "min_loc_n": min_loc_n})


def cell_log_index(res: SharpResult, level: str, geo_id: str, mode: str = "smoothed") -> tuple[np.ndarray, np.ndarray]:
    """Log index of a cell = provincial path + sum of its ancestors' and own deviations; variance = sum of state variances
    (independence across levels is an approximation; stage-1 μ uncertainty is added)."""
    path = res.stage1.mu.copy(); var = res.stage1.mu_se ** 2
    chain = []
    lvl_names = [l for l, _ in LEVELS]
    cur = (level, geo_id)
    while cur is not None and cur in res.cells:
        c = res.cells[cur]; chain.append(c)
        li = lvl_names.index(c.level)
        cur = (lvl_names[li - 1], c.parent_id) if li > 0 and c.parent_id is not None else None
    for c in chain:
        f = c.fit
        path = path + (f.smoothed if mode == "smoothed" else f.filtered)
        var = var + (f.smoothed_var if mode == "smoothed" else f.filtered_var)
    return path, var