# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Heteroskedastic local-level / local-linear-trend state-space models for cell deviations,
with robust (Huber-weighted innovation) filtering and empirical-Bayes pooling of innovation variances.

Observation : ȳ_t = d_t + ε_t,         ε_t ~ N(0, σ²/n_t)   (n_t = 0 → no observation)
Local level : d_t = d_{t-1} + η_t,     η_t ~ N(0, τ²)
Local trend : d_t = d_{t-1} + b_{t-1} + η_t ; b_t = b_{t-1} + ζ_t, ζ_t ~ N(0, τ_b²)

Outputs: filtered (one-sided, real-time) and smoothed (two-sided) paths and variances, the Kalman gain
(share of the period's estimate coming from local data) and the marginal log-likelihood.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np

LOG_TAU2_GRID = np.linspace(np.log(1e-8), np.log(2e-2), 48)
HUBER_K = 1.345


@dataclass
class SSFit:
    filtered: np.ndarray
    filtered_var: np.ndarray
    smoothed: np.ndarray
    smoothed_var: np.ndarray
    gain: np.ndarray
    tau2: float
    tau2_slope: float
    sigma2: float
    loglik: float
    model: str
    innov_weights: np.ndarray


def _filter_ll(y, r, tau2, p0):
    T = len(y)
    a = np.empty(T); p = np.empty(T); ap = np.empty(T); pp = np.empty(T); gain = np.zeros(T); v_std = np.zeros(T)
    at, pt, ll = 0.0, p0, 0.0
    for t in range(T):
        at_p, pt_p = at, pt + tau2
        ap[t], pp[t] = at_p, pt_p
        if np.isfinite(r[t]):
            f = pt_p + r[t]; v = y[t] - at_p; k = pt_p / f
            at = at_p + k * v; pt = pt_p * (1 - k); gain[t] = k; v_std[t] = v / np.sqrt(f)
            ll += -0.5 * (np.log(2 * np.pi * f) + v * v / f)
        else:
            at, pt = at_p, pt_p
        a[t], p[t] = at, pt
    return a, p, ap, pp, gain, ll, v_std


def _smooth_ll(a, p, ap, pp):
    T = len(a); s = a.copy(); sv = p.copy()
    for t in range(T - 2, -1, -1):
        c = p[t] / pp[t + 1] if pp[t + 1] > 0 else 0.0
        s[t] = a[t] + c * (s[t + 1] - ap[t + 1]); sv[t] = p[t] + c * c * (sv[t + 1] - pp[t + 1])
    return s, np.maximum(sv, 0)


def _filter_llt(y, r, tau2, tau2b, p0):
    """Local linear trend, 2-state Kalman filter + RTS smoother (returns level quantities)."""
    T = len(y)
    Tm = np.array([[1.0, 1.0], [0.0, 1.0]]); Q = np.diag([tau2, tau2b]); Z = np.array([1.0, 0.0])
    a = np.zeros((T, 2)); P = np.zeros((T, 2, 2)); ap = np.zeros((T, 2)); Pp = np.zeros((T, 2, 2)); gain = np.zeros(T); v_std = np.zeros(T)
    at = np.zeros(2); Pt = np.diag([p0, p0 / 10]); ll = 0.0
    for t in range(T):
        at_p = Tm @ at; Pt_p = Tm @ Pt @ Tm.T + Q
        ap[t], Pp[t] = at_p, Pt_p
        if np.isfinite(r[t]):
            f = Z @ Pt_p @ Z + r[t]; v = y[t] - Z @ at_p; K = Pt_p @ Z / f
            at = at_p + K * v; Pt = Pt_p - np.outer(K, Z @ Pt_p); gain[t] = K[0]; v_std[t] = v / np.sqrt(f)
            ll += -0.5 * (np.log(2 * np.pi * f) + v * v / f)
        else:
            at, Pt = at_p, Pt_p
        a[t], P[t] = at, Pt
    s = a.copy(); S = P.copy()
    for t in range(T - 2, -1, -1):
        try:
            C = P[t] @ Tm.T @ np.linalg.inv(Pp[t + 1])
        except np.linalg.LinAlgError:
            C = np.zeros((2, 2))
        s[t] = a[t] + C @ (s[t + 1] - ap[t + 1]); S[t] = P[t] + C @ (S[t + 1] - Pp[t + 1]) @ C.T
    return a[:, 0], P[:, 0, 0], s[:, 0], np.maximum(S[:, 0, 0], 0), gain, ll, v_std


def fit_cell(y: np.ndarray, n: np.ndarray, sigma2: float, model: str = "ll", log_tau2: float | None = None,
             prior: tuple[float, float] | None = None, robust: bool = True, p0: float = 0.25, tau2_slope_ratio: float = 0.05) -> SSFit:
    """Fit one cell. If `log_tau2` is None it is chosen on LOG_TAU2_GRID by maximising
    loglik + log N(log τ²; prior) (prior=None → maximum likelihood). Robust: two passes of Huber
    down-weighting of standardised innovations (observation variance inflated for outlying months)."""
    y = np.asarray(y, float); n = np.asarray(n, float)
    r0 = np.where(n > 0, sigma2 / np.maximum(n, 1), np.inf)
    yy = np.where(np.isfinite(r0), np.nan_to_num(y), 0.0)
    w = np.ones(len(y))

    def run(lt2, r):
        t2 = np.exp(lt2)
        if model == "llt":
            a, p, s, sv, g, ll, vs = _filter_llt(yy, r, t2, t2 * tau2_slope_ratio, p0)
        else:
            a, p, ap, pp, g, ll, vs = _filter_ll(yy, r, t2, p0); s, sv = _smooth_ll(a, p, ap, pp)
        return a, p, s, sv, g, ll, vs

    def choose(r):
        if log_tau2 is not None:
            return log_tau2
        lls = np.array([run(lt, r)[5] for lt in LOG_TAU2_GRID])
        if prior is not None:
            m, v = prior; lls = lls - 0.5 * (LOG_TAU2_GRID - m) ** 2 / v
        return float(LOG_TAU2_GRID[int(np.argmax(lls))])

    r = r0.copy(); lt = choose(r)
    if robust:
        for _ in range(2):
            vs = run(lt, r)[6]
            u = np.abs(vs); w = np.where(u <= HUBER_K, 1.0, HUBER_K / np.maximum(u, 1e-9)); w[~np.isfinite(r0)] = 1.0
            r = np.where(np.isfinite(r0), r0 / w, np.inf); lt = choose(r)
    a, p, s, sv, g, ll, vs = run(lt, r)
    t2 = float(np.exp(lt))
    return SSFit(a, p, s, sv, g, t2, t2 * tau2_slope_ratio if model == "llt" else 0.0, sigma2, ll, model, w)


def profile_loglik(y: np.ndarray, n: np.ndarray, sigma2: float, model: str = "ll", p0: float = 0.25) -> np.ndarray:
    """Log-likelihood of a cell on LOG_TAU2_GRID (for empirical-Bayes pooling across cells)."""
    y = np.asarray(y, float); n = np.asarray(n, float)
    r = np.where(n > 0, sigma2 / np.maximum(n, 1), np.inf); yy = np.where(np.isfinite(r), np.nan_to_num(y), 0.0)
    out = np.empty(len(LOG_TAU2_GRID))
    for i, lt in enumerate(LOG_TAU2_GRID):
        t2 = np.exp(lt)
        out[i] = _filter_llt(yy, r, t2, t2 * 0.05, p0)[5] if model == "llt" else _filter_ll(yy, r, t2, p0)[5]
    return out


def eb_prior(profiles: list[np.ndarray], n_iter: int = 50, v_floor: float = 0.05) -> tuple[float, float]:
    """Grid EM for a normal prior on log τ² shared by a group of cells: maximise Σ_c log ∫ L_c(τ²) N(log τ²; m, v)."""
    m, v = float(np.mean([LOG_TAU2_GRID[int(np.argmax(p))] for p in profiles])), 1.0
    g = LOG_TAU2_GRID
    for _ in range(n_iter):
        post_m, post_v = [], []
        for p in profiles:
            lp = p - 0.5 * (g - m) ** 2 / v; lp -= lp.max(); w = np.exp(lp); w /= w.sum()
            mu = np.sum(w * g); post_m.append(mu); post_v.append(np.sum(w * (g - mu) ** 2))
        m_new = float(np.mean(post_m)); v_new = float(max(np.mean(post_v) + np.var(post_m), v_floor))
        if abs(m_new - m) < 1e-6 and abs(v_new - v) < 1e-6:
            break
        m, v = m_new, v_new
    return m, v