# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task B models: forecast the h-month change of a cell's quality-adjusted log index from information at the origin.

Interface: fit(panel_train) ; predict(panel_test) → point forecast of `target`; predict_quantiles(panel_test, taus) optional.
`panel_*` are polars frames from data_b.build_panel restricted to one horizon h (train rows: origins whose target is observed
before the test origin — enforced by the caller).
Simple benchmarks: no-change (0), mean growth (historical mean h-change of the cell), momentum (last 3-m change scaled),
seasonal naive (same-calendar change last year), random walk with drift (12-m drift × h/12).
Time-series benchmarks on each cell's own history: AR(2) on monthly changes, local linear trend (Kalman).
Pooled panel benchmarks: ridge, LightGBM, MLP, LSTM / GRU / TCN on the 13-month history, TFT-lite (variable-selection + attention).
Proposed: DSGFN (graph attention over cells), RAREF (HMM regimes + gated experts), BHPF (hierarchical Bayesian ridge, posterior predictive),
QANSS (Kalman latent state per cell with learned nonlinear drift), QREFA (integration)."""
from __future__ import annotations

import os

import numpy as np
import polars as pl
from sklearn.linear_model import Ridge

from wp23.data_b import FEATURES_B

TAUS = np.array([0.05, 0.25, 0.5, 0.75, 0.95])


def _X(df: pl.DataFrame) -> np.ndarray:
    return np.nan_to_num(df.select(FEATURES_B).to_numpy().astype(float))


def _keys(df: pl.DataFrame, cols: list[str]) -> np.ndarray:
    return df.select(pl.concat_str([pl.col(c).cast(pl.Utf8).fill_null("na") for c in cols], separator="|").alias("k"))["k"].to_numpy().astype(str)


def _H(df: pl.DataFrame) -> np.ndarray:
    H = np.array(df["hist"].to_list(), float); return H - H[:, [-1]]  # history relative to the origin level


class Base:
    name = "base"; prob = False
    def fit(self, tr): return self
    def predict(self, te): raise NotImplementedError
    def predict_quantiles(self, te, taus=TAUS):
        """Default: empirical residual quantiles from training (homoskedastic)."""
        p = self.predict(te); return p[:, None] + np.quantile(self._res, taus)[None, :]
    def _store_res(self, tr):
        self._res = tr["target"].to_numpy() - self.predict(tr); return self


class NoChange(Base):
    name = "no_change"
    def fit(self, tr): return self._store_res(tr)
    def predict(self, te): return np.zeros(te.height)


class MeanGrowth(Base):
    name = "mean_growth"
    def fit(self, tr):
        self.h = int(tr["h"][0]); m = tr.group_by(["class", "level", "geo_id"]).agg(pl.col("target").mean().alias("mg")); self.map = {tuple(k): v for *k, v in m.iter_rows()}
        self.glob = float(tr["target"].mean()); return self._store_res(tr)
    def predict(self, te): return np.array([self.map.get((c, l, g), self.glob) for c, l, g in te.select(["class", "level", "geo_id"]).iter_rows()])


class Momentum(Base):
    name = "momentum"
    def fit(self, tr):
        self.h = int(tr["h"][0]); x = tr["d3"].to_numpy() / 3 * self.h; y = tr["target"].to_numpy(); self.b = float(np.sum(x * y) / max(np.sum(x * x), 1e-12)); return self._store_res(tr)
    def predict(self, te): return self.b * te["d3"].to_numpy() / 3 * self.h


class RWDrift(Base):
    name = "rw_drift"
    def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)
    def predict(self, te): return te["d12"].to_numpy() / 12 * self.h


class SeasonalNaive(Base):
    name = "seasonal_naive"
    def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)
    def predict(self, te):
        H = np.array(te["hist"].to_list(), float); h = self.h
        # change over the same months one year earlier, when the 13-month history covers it (h ≤ 12 → months -12..-12+h)
        if h <= 12 and H.shape[1] >= 13:
            return H[:, h] - H[:, 0]
        return np.zeros(te.height)


class AR2(Base):
    name = "ar2_cell"
    def fit(self, tr):
        self.h = int(tr["h"][0]); D = np.diff(_H(tr), axis=1); y = tr["target"].to_numpy()
        X = np.column_stack([np.ones(len(D)), D[:, -1], D[:, -2]]); self.b = np.linalg.lstsq(X, y, rcond=None)[0]; return self._store_res(tr)
    def predict(self, te): D = np.diff(_H(te), axis=1); return np.column_stack([np.ones(len(D)), D[:, -1], D[:, -2]]) @ self.b


class LocalLinearTrend(Base):
    name = "local_linear_trend"
    def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)
    def predict(self, te):
        import sys
        from wp23 import config as C
        sys.path.insert(0, str(C.WP22_SRC)); from sharp_hpi.statespace import _filter_llt
        out = np.empty(te.height); H = np.array(te["hist"].to_list(), float)
        for i, y in enumerate(H):
            r = np.full(len(y), 1e-4); a, p, s, sv, g, ll, vs = _filter_llt(y, r, 1e-5, 2e-6, 1.0)
            # slope at the end from the filtered state (re-run with 2-state output): approximate slope = last filtered level change
            out[i] = (a[-1] - a[-2]) * self.h
        return out


class RidgePanel(Base):
    name = "ridge_panel"
    def fit(self, tr):
        X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; self.m = Ridge(alpha=3.0).fit((X - self.mu) / self.sd, tr["target"].to_numpy()); return self._store_res(tr)
    def predict(self, te): return self.m.predict((_X(te) - self.mu) / self.sd)


class LGBMPanel(Base):
    name = "lightgbm_panel"; prob = True
    def fit(self, tr):
        import lightgbm as lgb
        X = _X(tr); y = tr["target"].to_numpy(); kw = dict(n_estimators=300, learning_rate=0.03, num_leaves=15, min_child_samples=40, subsample=0.8, subsample_freq=1, colsample_bytree=0.8, verbose=-1, n_jobs=4, random_state=0)
        self.m = lgb.LGBMRegressor(objective="l2", **kw).fit(X, y); self.q = {t: lgb.LGBMRegressor(objective="quantile", alpha=float(t), **kw).fit(X, y) for t in TAUS}; return self._store_res(tr)
    def predict(self, te): return self.m.predict(_X(te))
    def predict_quantiles(self, te, taus=TAUS): return np.sort(np.column_stack([self.q[t].predict(_X(te)) for t in taus]), axis=1)


class GaussianHMM:
    """Two-state Gaussian HMM on a univariate series, EM (Baum–Welch); used on the provincial/regional real-time growth at the origin."""
    def __init__(self, K=3, iters=60, seed=0): self.K, self.iters, self.rng = K, iters, np.random.default_rng(seed)
    def fit(self, x):
        x = np.asarray(x, float); K = self.K; q = np.quantile(x, np.linspace(0.2, 0.8, K)); mu = q.copy(); sd = np.full(K, x.std() + 1e-6); A = np.full((K, K), 0.1) + np.eye(K) * (0.9 - 0.1 * (K - 1)); pi = np.full(K, 1 / K)
        for _ in range(self.iters):
            B = np.exp(-0.5 * ((x[:, None] - mu) / sd) ** 2) / sd; T = len(x); al = np.zeros((T, K)); be = np.zeros((T, K)); c = np.zeros(T)
            al[0] = pi * B[0]; c[0] = al[0].sum() + 1e-300; al[0] /= c[0]
            for t in range(1, T):
                al[t] = (al[t - 1] @ A) * B[t]; c[t] = al[t].sum() + 1e-300; al[t] /= c[t]
            be[-1] = 1
            for t in range(T - 2, -1, -1):
                be[t] = (A @ (B[t + 1] * be[t + 1])) / c[t + 1]
            g = al * be; g /= g.sum(1, keepdims=True) + 1e-300
            xi = np.zeros((K, K))
            for t in range(T - 1):
                m = al[t][:, None] * A * (B[t + 1] * be[t + 1])[None, :]; xi += m / (m.sum() + 1e-300)
            pi = g[0]; A = xi / (xi.sum(1, keepdims=True) + 1e-300); mu = (g * x[:, None]).sum(0) / (g.sum(0) + 1e-300); sd = np.sqrt((g * (x[:, None] - mu) ** 2).sum(0) / (g.sum(0) + 1e-300)) + 1e-4
        self.mu, self.sd, self.A, self.pi = mu, sd, A, pi; self.filtered = al; return self
    def filter_last(self, x):
        x = np.asarray(x, float); B = np.exp(-0.5 * ((x[:, None] - self.mu) / self.sd) ** 2) / self.sd; al = self.pi * B[0]; al /= al.sum() + 1e-300
        for t in range(1, len(x)):
            al = (al @ self.A) * B[t]; al /= al.sum() + 1e-300
        return al


class RAREF(Base):
    """Regime-adaptive forecasting: a 3-state Gaussian HMM is estimated at each origin on the provincial real-time monthly growth
    (information ≤ origin only); the filtered regime probabilities gate a mixture of ridge experts (one per regime) estimated on
    training rows weighted by their own origin's regime probabilities. Forecast = probability-weighted expert forecast."""
    name = "raref"; prob = True
    def __init__(self, K=3): self.K = K
    def _probs(self, df):
        # per origin: filtered regime probabilities of the provincial growth series up to that origin
        out = np.zeros((df.height, self.K)); cache = {}
        for i, (o, c) in enumerate(df.select(["origin", "class"]).iter_rows()):
            if (o, c) not in cache:
                hist = df.filter((pl.col("origin") == o) & (pl.col("class") == c) & (pl.col("level") == "province"))
                if hist.height == 0:
                    cache[(o, c)] = np.full(self.K, 1 / self.K)
                else:
                    g = np.diff(np.array(hist["hist"][0], float)); key = (o, c)
                    series = self._growth_series(df, o, c)
                    cache[key] = self.hmm_for(series).filter_last(series) if len(series) > 10 else np.full(self.K, 1 / self.K)
            out[i] = cache[(o, c)]
        return out
    def _growth_series(self, df, o, c):
        # concatenate provincial growth across origins ≤ o from the stored histories (real-time at each origin): use the origin's own 13-month history
        # extended backward with the earliest available origins (all ≤ o → no leakage)
        hs = df.filter((pl.col("class") == c) & (pl.col("level") == "province") & (pl.col("origin") <= o)).sort("origin")
        if hs.height == 0:
            return np.array([])
        seq = []
        for o2, hh in zip(hs["origin"].to_list(), hs["hist"].to_list()):
            seq.append(np.diff(np.array(hh, float))[-1])  # last monthly growth known at each past origin
        return np.array(seq)
    def hmm_for(self, series):
        return GaussianHMM(self.K).fit(series)
    def fit(self, tr):
        self.P = self._probs(tr); X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; Xs = (X - self.mu) / self.sd; y = tr["target"].to_numpy()
        self.experts = []
        for k in range(self.K):
            w = self.P[:, k] + 1e-3; self.experts.append(Ridge(alpha=3.0).fit(Xs, y, sample_weight=w))
        self._train_df = tr; return self._store_res(tr)
    def predict(self, te):
        both = pl.concat([self._train_df.select(te.columns), te], how="diagonal") if te.height else te
        P = self._probs(both)[-te.height:]; Xs = (_X(te) - self.mu) / self.sd
        return sum(P[:, k] * self.experts[k].predict(Xs) for k in range(self.K))


class BHPF(Base):
    """Bayesian hierarchical probabilistic forecasting: hierarchical linear model with cell-level random intercepts (shrunk toward
    their region, then the class mean) and common slopes on the standardized features; conjugate Gaussian posterior (ridge-type)
    with variance components estimated by REML-like iteration; predictive distribution = posterior predictive Gaussian with
    cell-specific variance. Laplace/conjugate rather than MCMC: exact for this Gaussian model."""
    name = "bhpf"; prob = True
    def fit(self, tr):
        X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; Xs = (X - self.mu) / self.sd; y = tr["target"].to_numpy()
        cells = _keys(tr, ["class", "level", "geo_id"]); regs = _keys(tr, ["class", "region_code"])
        self.cells = sorted(set(cells)); self.regs = sorted(set(regs)); ci = {c: i for i, c in enumerate(self.cells)}; ri = {r: i for i, r in enumerate(self.regs)}
        Zc = np.zeros((len(y), len(self.cells))); Zc[np.arange(len(y)), [ci[c] for c in cells]] = 1
        Zr = np.zeros((len(y), len(self.regs))); Zr[np.arange(len(y)), [ri[r] for r in regs]] = 1
        self.ci, self.ri = ci, ri
        D = np.column_stack([np.ones(len(y)), Xs, Zr, Zc]); p = Xs.shape[1]
        tau_b, tau_r, tau_c, s2 = 1.0, 0.01, 0.01, float(y.var())
        for _ in range(15):
            prec = np.concatenate([[1e-6], np.full(p, 1 / tau_b), np.full(len(self.regs), 1 / tau_r), np.full(len(self.cells), 1 / tau_c)])
            Ainv = np.linalg.inv(D.T @ D / s2 + np.diag(prec)); beta = Ainv @ (D.T @ y / s2); r = y - D @ beta
            s2 = float(np.mean(r ** 2)) + 1e-9
            br = beta[1 + p: 1 + p + len(self.regs)]; bc = beta[1 + p + len(self.regs):]
            tau_r = max(float(np.mean(br ** 2) + np.mean(np.diag(Ainv)[1 + p: 1 + p + len(self.regs)])), 1e-6); tau_c = max(float(np.mean(bc ** 2) + np.mean(np.diag(Ainv)[1 + p + len(self.regs):])), 1e-6)
        self.beta, self.Ainv, self.s2, self.p = beta, Ainv, s2, p; self.vc = {"s2": s2, "tau_region": tau_r, "tau_cell": tau_c}; return self._store_res(tr)
    def _design(self, te):
        Xs = (_X(te) - self.mu) / self.sd; cells = _keys(te, ["class", "level", "geo_id"]); regs = _keys(te, ["class", "region_code"])
        Zc = np.zeros((te.height, len(self.cells))); Zr = np.zeros((te.height, len(self.regs)))
        for i, (c, r) in enumerate(zip(cells, regs)):
            if c in self.ci: Zc[i, self.ci[c]] = 1
            if r in self.ri: Zr[i, self.ri[r]] = 1
        return np.column_stack([np.ones(te.height), Xs, Zr, Zc])
    def predict(self, te): return self._design(te) @ self.beta
    def predict_quantiles(self, te, taus=TAUS):
        from scipy.stats import norm
        D = self._design(te); m = D @ self.beta; v = np.einsum("ij,jk,ik->i", D, self.Ainv, D) + self.s2
        return m[:, None] + np.sqrt(v)[:, None] * norm.ppf(taus)[None, :]


def registry(family: str = "classic", fast: bool = False) -> list:
    if family == "torch":
        from wp23.models_b_torch import DSGFN, GRUPanel, LSTMPanel, MLPPanel, QANSS, QREFA, TCNPanel, TFTLite
        base = [MLPPanel(), LSTMPanel(), GRUPanel(), TCNPanel(), TFTLite(), DSGFN(), QANSS()]
        if not fast:
            base.append(QREFA())
        return base
    if family == "trees":
        return [LGBMPanel()]
    return [NoChange(), MeanGrowth(), Momentum(), RWDrift(), SeasonalNaive(), AR2(), LocalLinearTrend(), RidgePanel(), RAREF(), BHPF()]