# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Benchmark index estimators (protocol §5, A–J) on one geography × class sample.

Every estimator returns an `IndexSeries` of **log** index levels on a continuous month grid
(NaN where not estimable). Normalisation (2021 average = 100 or January 2021 = 100) is applied
downstream in `panel.py`, so it never changes growth rates.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import polars as pl
import pyfixest as pf
from scipy import sparse
from scipy.sparse.linalg import lsqr

from sharp_hpi.hedonic import RHS_STRUCT, fit_fe, huber_weights


@dataclass
class IndexSeries:
    months: list[str]
    log_index: np.ndarray
    se: np.ndarray | None = None
    n: np.ndarray | None = None
    method: str = ""
    meta: dict = field(default_factory=dict)

    def to_frame(self) -> pl.DataFrame:
        d = {"month": self.months, "log_index": self.log_index, "method": self.method}
        if self.se is not None:
            d["se"] = self.se
        if self.n is not None:
            d["n"] = self.n
        return pl.DataFrame(d)


def month_grid(first: str, last: str) -> list[str]:
    y, m = int(first[:4]), int(first[5:7])
    out = []
    while True:
        cur = f"{y:04d}-{m:02d}"
        out.append(cur)
        if cur == last:
            return out
        m += 1
        if m == 13:
            y, m = y + 1, 1


def _counts(df: pl.DataFrame, grid: list[str]) -> np.ndarray:
    c = df.group_by("month").agg(pl.len().alias("n"))
    d = dict(zip(c["month"].to_list(), c["n"].to_list()))
    return np.array([d.get(m, 0) for m in grid], float)


def _rhs_for(df: pl.DataFrame) -> str:
    rhs = RHS_STRUCT
    if df["property_class"].n_unique() == 1 and df["property_class"][0] != "plex":
        rhs = rhs.replace(" + C(units_c)", "")
    if df["property_class"].n_unique() == 1 and df["property_class"][0] != "house":
        rhs = rhs.replace(" + C(subtype_c)", "")
    return rhs


# A. raw median / mean -------------------------------------------------------
def median_index(df: pl.DataFrame, grid: list[str], stat: str = "median") -> IndexSeries:
    agg = pl.col("price").median() if stat == "median" else pl.col("price").mean()
    g = df.group_by("month").agg(agg.alias("v")).sort("month")
    d = dict(zip(g["month"].to_list(), g["v"].to_list()))
    li = np.array([np.log(d[m]) if m in d else np.nan for m in grid])
    return IndexSeries(grid, li, n=_counts(df, grid), method=f"A_{stat}")


# B. stratified median (fixed-weight Laspeyres over strata) ------------------
def stratified_median(df: pl.DataFrame, grid: list[str], n_strata_fa: int = 3, n_strata_age: int = 3) -> IndexSeries:
    qfa = [df["fa_f"].quantile(q) for q in np.linspace(0, 1, n_strata_fa + 1)[1:-1]]
    qage = [df["age_f"].quantile(q) for q in np.linspace(0, 1, n_strata_age + 1)[1:-1]]
    d = df.with_columns(
        pl.col("fa_f").cut(qfa, labels=[f"fa{i}" for i in range(n_strata_fa)]).cast(pl.Utf8).alias("_sfa"),
        pl.col("age_f").cut(qage, labels=[f"ag{i}" for i in range(n_strata_age)]).cast(pl.Utf8).alias("_sag"),
    ).with_columns(pl.concat_str([pl.col("_sfa"), pl.col("_sag"), pl.col("subtype_c")]).alias("stratum"))
    w = d.group_by("stratum").agg(pl.len().alias("w")).with_columns(pl.col("w") / pl.col("w").sum())
    med = d.group_by(["stratum", "month"]).agg(pl.col("price").median().log().alias("lm"))
    strata = w["stratum"].to_list()
    wv = dict(zip(strata, w["w"].to_list()))
    lookup = {(s, m): v for s, m, v in zip(med["stratum"].to_list(), med["month"].to_list(), med["lm"].to_list())}
    out = np.full(len(grid), np.nan)
    last = {s: np.nan for s in strata}
    for i, m in enumerate(grid):
        acc, wsum = 0.0, 0.0
        for s in strata:
            v = lookup.get((s, m), np.nan)
            if np.isnan(v):
                v = last[s]  # carry forward an empty stratum's last median
            else:
                last[s] = v
            if not np.isnan(v):
                acc += wv[s] * v
                wsum += wv[s]
        out[i] = acc / wsum if wsum > 0.5 else np.nan
    return IndexSeries(grid, out, n=_counts(df, grid), method="B_stratified_median", meta={"n_strata": len(strata)})


# C. pooled hedonic time dummy (Huber-IRLS, location FE) ----------------------
def pooled_td(df: pl.DataFrame, grid: list[str], robust: bool = True, fe: str = "loc_fe") -> IndexSeries:
    rhs = _rhs_for(df) + " + C(month)"
    res, m = fit_fe(df, rhs, fe, robust=robust)
    coef = res.coef
    se = dict(m.se())
    base = grid[0]
    li = np.array([0.0 if g == base else coef.get(f"C(month)[T.{g}]", np.nan) for g in grid])
    s = np.array([0.0 if g == base else se.get(f"C(month)[T.{g}]", np.nan) for g in grid])
    return IndexSeries(grid, li, se=s, n=_counts(df, grid), method="C_pooled_td" + ("" if robust else "_ols"),
                       meta={"sigma_robust": res.sigma2_robust ** 0.5, "n": res.n})


# D. rolling time dummy with window / movement / mean splice ------------------
def rolling_td(df: pl.DataFrame, grid: list[str], window: int = 13, splice: str = "mean", robust: bool = True, fe: str = "loc_fe",
               min_obs: int = 50) -> IndexSeries:
    """Hill et al.-type rolling time dummy. Each window of `window` months is estimated; the published path is
    extended one month at a time. splice ∈ {window, movement, mean, half}."""
    rhs = _rhs_for(df) + " + C(month)"
    T = len(grid)
    path = np.full(T, np.nan)
    paths_w: list[np.ndarray] = []
    first_end = min(window, T)
    for end in range(first_end, T + 1):
        w_months = grid[end - window: end] if end >= window else grid[:end]
        sub = df.filter(pl.col("month").is_in(w_months))
        if sub.height < min_obs or sub["month"].n_unique() < 2:
            paths_w.append(np.full(len(w_months), np.nan)); continue
        res, m = fit_fe(sub, rhs, fe, robust=robust)
        coef = res.coef
        b = w_months[0]
        pw = np.array([0.0 if g == b else coef.get(f"C(month)[T.{g}]", np.nan) for g in w_months])
        paths_w.append(pw)
        if end == first_end:
            path[:end] = pw
            continue
        new_pos = end - 1
        if splice == "window":
            path[new_pos] = path[end - window] + (pw[-1] - pw[0])
        elif splice == "movement":
            path[new_pos] = path[new_pos - 1] + (pw[-1] - pw[-2])
        elif splice == "half":
            k = max(1, (window - 1) // 2)
            path[new_pos] = path[new_pos - k] + (pw[-1] - pw[-1 - k])
        else:  # mean splice (Diewert–Fox / Hill et al.): average over all overlap anchors
            cands = []
            for j in range(window - 1):
                anchor = new_pos - (window - 1) + j
                if anchor >= 0 and not np.isnan(path[anchor]) and not np.isnan(pw[j]):
                    cands.append(path[anchor] + (pw[-1] - pw[j]))
            path[new_pos] = np.mean(cands) if cands else np.nan
    return IndexSeries(grid, path, n=_counts(df, grid), method=f"D_rtd_w{window}_{splice}")


# E/F. hedonic imputation (Laspeyres / Paasche / Fisher) and GEKS ------------
def _period_fits(df: pl.DataFrame, grid: list[str], loc_col: str, min_obs: int = 40, robust: bool = False):
    """Per-month hedonic fits (OLS or Huber) with a coarse location FE; returns {month: (model, pandas data)}."""
    rhs = _rhs_for(df)
    fits = {}
    for mth in grid:
        sub = df.filter(pl.col("month") == mth)
        if sub.height < min_obs:
            continue
        try:
            res, m = fit_fe(sub, rhs, loc_col, robust=robust, n_irls=2)
        except Exception:
            continue
        fits[mth] = (m, sub)
    return fits


def _impute(m, data_pl: pl.DataFrame) -> np.ndarray:
    pdf = data_pl.to_pandas()
    try:
        pred = np.asarray(m.predict(newdata=pdf), float)
    except Exception:
        return np.full(len(pdf), np.nan)
    return pred


def hedonic_imputation(df: pl.DataFrame, grid: list[str], loc_col: str = "loc_coarse", kind: str = "fisher", chained: bool = True,
                       min_obs: int = 40, robust: bool = False) -> IndexSeries:
    fits = _period_fits(df, grid, loc_col, min_obs, robust)
    T = len(grid)
    li = np.full(T, np.nan)
    if not fits:
        return IndexSeries(grid, li, n=_counts(df, grid), method=f"E_hi_{kind}")
    months = [g for g in grid if g in fits]
    li[grid.index(months[0])] = 0.0

    def bilateral(s: str, t: str) -> float:
        ms, ds = fits[s]; mt, dt = fits[t]
        # Laspeyres: basket s, prices t vs s (double imputation)
        ps_s = _impute(ms, ds); pt_s = _impute(mt, ds)
        ok = ~np.isnan(ps_s) & ~np.isnan(pt_s)
        lasp = np.mean(pt_s[ok] - ps_s[ok]) if ok.sum() >= 10 else np.nan
        ps_t = _impute(ms, dt); pt_t = _impute(mt, dt)
        ok2 = ~np.isnan(ps_t) & ~np.isnan(pt_t)
        paas = np.mean(pt_t[ok2] - ps_t[ok2]) if ok2.sum() >= 10 else np.nan
        return {"laspeyres": lasp, "paasche": paas, "fisher": 0.5 * (lasp + paas)}[kind]

    if chained:
        for a, b in zip(months[:-1], months[1:]):
            li[grid.index(b)] = li[grid.index(a)] + bilateral(a, b)
    else:
        for b in months[1:]:
            li[grid.index(b)] = bilateral(months[0], b)
    return IndexSeries(grid, li, n=_counts(df, grid), method=f"E_hi_{kind}{'_chained' if chained else '_direct'}", meta={"months_estimated": len(months)})


def geks(df: pl.DataFrame, grid: list[str], loc_col: str = "loc_coarse", window: int = 25, min_obs: int = 40, robust: bool = False) -> IndexSeries:
    """Rolling-window GEKS on bilateral double-imputation Fisher links, extended by movement splice (RYGEKS)."""
    fits = _period_fits(df, grid, loc_col, min_obs, robust)
    months = [g for g in grid if g in fits]
    T = len(months)
    li = np.full(len(grid), np.nan)
    if T < 2:
        return IndexSeries(grid, li, n=_counts(df, grid), method=f"F_geks_w{window}")
    # cache imputed prices: pred[t][s] = f_t applied to basket s
    pred: dict[tuple[str, str], np.ndarray] = {}

    def P(t: str, s: str) -> np.ndarray:
        if (t, s) not in pred:
            pred[(t, s)] = _impute(fits[t][0], fits[s][1])
        return pred[(t, s)]

    def fisher(s: str, t: str) -> float:
        ok = ~np.isnan(P(t, s)) & ~np.isnan(P(s, s)); lasp = np.mean((P(t, s) - P(s, s))[ok]) if ok.sum() >= 10 else np.nan
        ok2 = ~np.isnan(P(t, t)) & ~np.isnan(P(s, t)); paas = np.mean((P(t, t) - P(s, t))[ok2]) if ok2.sum() >= 10 else np.nan
        return 0.5 * (lasp + paas)

    def geks_window(ws: list[str]) -> np.ndarray:
        k = len(ws)
        F = np.zeros((k, k))
        for i in range(k):
            for j in range(i + 1, k):
                F[i, j] = fisher(ws[i], ws[j]); F[j, i] = -F[i, j]
        F = np.where(np.isnan(F), 0.0, F)
        return F.mean(axis=0)  # mean_j F[j,t] = log P_t − mean_j log P_j (GEKS, up to a constant)
    out = np.full(T, np.nan)
    first = min(window, T)
    g0 = geks_window(months[:first]); out[:first] = g0 - g0[0]
    for end in range(first + 1, T + 1):
        ws = months[end - window: end]
        g = geks_window(ws)
        out[end - 1] = out[end - 2] + (g[-1] - g[-2])  # movement splice
    for i, m in enumerate(months):
        li[grid.index(m)] = out[i]
    return IndexSeries(grid, li, n=_counts(df, grid), method=f"F_geks_w{window}_movement")


# H. repeat sales -------------------------------------------------------------
def repeat_sales_pairs(df: pl.DataFrame, min_hold_months: int = 6) -> pl.DataFrame:
    d = df.filter(pl.col("parcel_id").is_not_null()).sort(["parcel_id", "date"]).with_columns(
        pl.col("month").shift(1).over("parcel_id").alias("m1"), pl.col("log_price").shift(1).over("parcel_id").alias("lp1"),
        pl.col("date").shift(1).over("parcel_id").alias("d1"),
    ).filter(pl.col("m1").is_not_null())
    d = d.with_columns(((pl.col("date") - pl.col("d1")).dt.total_days() / 30.4375).alias("hold_months"))
    return d.filter(pl.col("hold_months") >= min_hold_months).select(["parcel_id", "m1", pl.col("month").alias("m2"), "lp1", pl.col("log_price").alias("lp2"), "hold_months"])


def repeat_sales(df: pl.DataFrame, grid: list[str], weighted: bool = True, min_hold_months: int = 6) -> IndexSeries:
    pairs = repeat_sales_pairs(df, min_hold_months)
    T = len(grid)
    li = np.full(T, np.nan)
    if pairs.height < 30:
        return IndexSeries(grid, li, method="H_repeat_sales")
    idx = {m: i for i, m in enumerate(grid)}
    i1 = np.array([idx[m] for m in pairs["m1"].to_list()]); i2 = np.array([idx[m] for m in pairs["m2"].to_list()])
    y = (pairs["lp2"] - pairs["lp1"]).to_numpy()
    n = len(y)
    rows = np.concatenate([np.arange(n), np.arange(n)]); cols = np.concatenate([i1, i2]); vals = np.concatenate([-np.ones(n), np.ones(n)])
    X = sparse.csr_matrix((vals, (rows, cols)), shape=(n, T))[:, 1:]  # base = first month
    beta = lsqr(X, y, atol=1e-10, btol=1e-10)[0]
    if weighted:  # Case–Shiller: var(resid) = a + b*hold
        r = y - X @ beta
        h = pairs["hold_months"].to_numpy()
        A = np.column_stack([np.ones(n), h]); ab = np.linalg.lstsq(A, r ** 2, rcond=None)[0]
        v = np.maximum(A @ ab, 1e-4); w = 1 / np.sqrt(v)
        beta = lsqr(sparse.diags(w) @ X, w * y, atol=1e-10, btol=1e-10)[0]
    li[0] = 0.0; li[1:] = beta
    cnt = np.bincount(np.concatenate([i1, i2]), minlength=T).astype(float)
    # months with no pair endpoints are not identified
    li[cnt == 0] = np.nan
    return IndexSeries(grid, li, n=cnt, method="H_repeat_sales" + ("_cs" if weighted else "_bmn"), meta={"pairs": n})


# I. static empirical-Bayes partial pooling of cell deviations (mixed-effects benchmark) ---------
def static_eb_deviation(resid_mean: np.ndarray, n: np.ndarray, sigma2: float) -> tuple[np.ndarray, np.ndarray, float]:
    """Random-intercept shrinkage per period: deviation_t ~ N(0, tau2); ybar_t | dev ~ N(dev, sigma2/n_t).
    tau2 by marginal likelihood (method of moments fallback). Returns (posterior mean, weight, tau2)."""
    ok = n > 0
    if ok.sum() < 2:
        return np.zeros_like(resid_mean), np.zeros_like(resid_mean), 0.0
    y = resid_mean[ok]; v = sigma2 / n[ok]
    tau2 = max(np.var(y) - np.mean(v), 1e-8)
    for _ in range(50):  # DerSimonian–Laird style iteration
        w = 1 / (v + tau2); mu = np.sum(w * y) / np.sum(w)
        tau2_new = max(np.sum(w ** 2 * ((y - mu) ** 2 - v)) / np.sum(w ** 2), 1e-8)
        if abs(tau2_new - tau2) < 1e-10:
            break
        tau2 = tau2_new
    wgt = np.zeros_like(resid_mean); post = np.zeros_like(resid_mean)
    wgt[ok] = tau2 / (tau2 + v); post[ok] = wgt[ok] * y
    return post, wgt, tau2


# J. out-of-time gradient boosting residual index (diagnostic only) ------------------------------
def gbm_residual_index(df: pl.DataFrame, grid: list[str], retrain_every: int = 3, min_train: int = 2000, seed: int = 20261021) -> IndexSeries:
    """Train LightGBM on log price ~ characteristics + coordinates using only months < t (expanding window,
    retrained every `retrain_every` months), then index_t = mean out-of-sample residual (chain of monthly
    quality-adjusted price levels relative to the training-period average). Predictive accuracy is NOT index accuracy."""
    import lightgbm as lgb
    feats = ["lfa", "lla", "age_f", "fa_missing", "la_missing", "age_missing", "x_m", "y_m"]
    cat = ["storeys_c", "genre_c", "subtype_c", "units_c"]
    d = df.with_columns([pl.col(c).cast(pl.Categorical).to_physical().alias(c + "_i") for c in cat])
    X_cols = feats + [c + "_i" for c in cat]
    T = len(grid)
    li = np.full(T, np.nan)
    model = None
    train_mean_month = None
    for i, mth in enumerate(grid):
        past = d.filter(pl.col("month") < mth)
        if past.height < min_train:
            continue
        if model is None or i % retrain_every == 0:
            # detrend training target with month means so the model learns quality, not time
            mm = past.group_by("month").agg(pl.col("log_price").mean().alias("_mm"))
            tr = past.join(mm, on="month")
            yt = (tr["log_price"] - tr["_mm"]).to_numpy()
            model = lgb.LGBMRegressor(n_estimators=400, learning_rate=0.05, num_leaves=63, min_child_samples=20, subsample=0.8, subsample_freq=1,
                                      colsample_bytree=0.8, objective="l1", random_state=seed, verbose=-1)
            model.fit(tr.select(X_cols).to_numpy(), yt)
            train_mean_month = mm
        cur = d.filter(pl.col("month") == mth)
        if cur.height == 0:
            continue
        pred = model.predict(cur.select(X_cols).to_numpy())
        li[i] = float(np.mean(cur["log_price"].to_numpy() - pred))
    return IndexSeries(grid, li, n=_counts(df, grid), method="J_gbm_oot_residual")


# G. HMTS — hedonic multilateral time-series (Ishaak, Ouwehand & Remøy 2024, JOS 40(3)) ---------------------
def hmts(df: pl.DataFrame, grid: list[str], loc_col: str = "loc_coarse", window: int = 13, min_obs: int = 40, start_periods: int = 20) -> IndexSeries:
    """For every base period b, the series p̂_{t(b)} = exp(α̂_t + β̂_t' x̄_b), t = b … b+w−1 (period-t hedonic coefficients on the
    base-period mean characteristics; here evaluated as the mean imputed log price on the base basket, eq. 4), is smoothed by a
    smooth-trend state-space model (local linear trend with zero level disturbance, ML hyperparameters, Kalman smoother; eqs. 9–12)
    fitted on log p̂; each smoothed series is indexed on its own first period (eq. 17) and the published index is the geometric mean
    over the w most recent base-period series (eq. 18). Periods beyond the first `start_periods` are 'limited-revisable' (w−1 preliminary
    periods). See docs/literature/hmts_method.md. Monthly adaptation: w = `window` (paper: w = 4 quarters)."""
    from scipy.optimize import minimize

    from sharp_hpi.statespace import _filter_llt
    fits = _period_fits(df, grid, loc_col, min_obs, robust=False)
    months = [g for g in grid if g in fits]
    T = len(months)
    li = np.full(len(grid), np.nan)
    if T < 3:
        return IndexSeries(grid, li, n=_counts(df, grid), method=f"G_hmts_w{window}")
    # imputed log prices of basket b at period-t prices (row vectors, cached); links use the common support of
    # consecutive periods so that rows unpredictable in one period (unseen location levels) do not shift the mean
    cache: dict[tuple[str, str], np.ndarray] = {}

    def P(t: str, b: str) -> np.ndarray:
        if (t, b) not in cache:
            cache[(t, b)] = _impute(fits[t][0], fits[b][1])
        return cache[(t, b)]

    def link(b: str, t0: str, t1: str) -> float:
        a0, a1 = P(t0, b), P(t1, b); ok = np.isfinite(a0) & np.isfinite(a1)
        return float(np.mean(a1[ok] - a0[ok])) if ok.sum() >= 10 else np.nan

    def smooth_trend(y: np.ndarray) -> np.ndarray:
        ok = ~np.isnan(y)
        if ok.sum() < 3:
            return y
        yy = np.where(ok, y, 0.0)

        def nll(th):
            th = np.clip(th, -18.0, 0.0)
            se2, sz2 = np.exp(th)
            r = np.where(ok, se2, np.inf)
            v = -_filter_llt(yy, r, 0.0, sz2, 1.0)[5]
            return v if np.isfinite(v) else 1e12
        best = minimize(nll, x0=[np.log(1e-3), np.log(1e-4)], method="Nelder-Mead", options={"maxiter": 300, "xatol": 1e-3, "fatol": 1e-4})
        se2, sz2 = np.exp(np.clip(best.x, -18.0, 0.0))
        s = _filter_llt(yy, np.where(ok, se2, np.inf), 0.0, sz2, 1.0)[2]
        if not np.all(np.isfinite(s[ok])) or np.nanmax(np.abs(s - np.nanmean(s))) > 1.0:
            return y  # degenerate fit on a very short / sparse series: keep the raw imputed series
        return s

    # base-period series: each extended to at most `window` periods, then frozen (Table 5)
    series: dict[int, np.ndarray] = {}
    for bi in range(T):
        end = min(bi + window, T)
        lk = [0.0] + [link(months[bi], months[t - 1], months[t]) for t in range(bi + 1, end)]
        y = np.cumsum(np.nan_to_num(lk)); y[np.isnan(lk)] = np.nan
        s = smooth_trend(y)
        series[bi] = s - s[0]  # indexed on its first period (eq. 17)
    # final index: for period t, geometric mean over the base-period series that cover t; chain across the moving set (eq. 18–19)
    out = np.full(T, np.nan); out[0] = 0.0
    for t in range(1, T):
        links = []
        for bi, s in series.items():
            if bi <= t - 1 and t - bi < len(s):
                links.append(s[t - bi] - s[t - 1 - bi])
        out[t] = out[t - 1] + (np.nanmean(links) if links else np.nan)
    for i, m in enumerate(months):
        li[grid.index(m)] = out[i]
    return IndexSeries(grid, li, n=_counts(df, grid), method=f"G_hmts_w{window}", meta={"months_estimated": T})