# QC Élection Forecast — Plateforme de prévision électorale du Québec 2026
# Auteur : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# https://www.qc-election.com
"""Modèle dynamique bayésien des intentions de vote (state space / Kalman).

État latent x_t ∈ R^5 : log-ratios des parts de vote (référence = CAQ).
Dynamique : marche aléatoire gaussienne x_t = x_{t-1} + w_t, w_t ~ N(0, q·Δt·I).
Observations : chaque sondage est une observation bruitée y = alr(parts ajustées)
avec covariance issue de la multinomiale (méthode delta) + erreur excédentaire.

La variance quotidienne q est estimée par maximum de vraisemblance (recherche
sur grille sur la log-vraisemblance des innovations), bornée par configuration.
"""
from __future__ import annotations

from dataclasses import dataclass
from datetime import date, timedelta

import numpy as np

from ..config import settings
from .compositions import alr, alr_obs_cov, close, inv_alr
from .house_effects import PollsterProfile, adjust_shares, profile_for
from .weights import effective_n

M = len(settings.parties) - 1  # dimension de l'état (K-1) — élection cible

# Le cœur est paramétré par la liste de partis : le replay historique rejoue
# des ères différentes (ADQ 2007-2008, pas de PCQ avant 2022…) avec le même code.


def shares_vector(shares: dict[str, float],
                  parties: list[str] | None = None) -> np.ndarray:
    """dict {parti: %} → vecteur ordonné (fractions, somme 1)."""
    ps = parties or settings.parties
    v = np.array([shares.get(p, 0.0) for p in ps], dtype=float)
    if "AUT" in ps and v[ps.index("AUT")] == 0.0:
        v[ps.index("AUT")] = max(0.5, 100.0 - v.sum())
    return close(v / 100.0)


@dataclass
class PreparedPoll:
    field_end: date
    pollster: str
    y: np.ndarray          # observation alr (5,)
    R: np.ndarray          # covariance d'observation (5,5)
    p_adj: np.ndarray      # parts ajustées (6,)
    n_eff: float
    multiplier: float


@dataclass
class TrendResult:
    dates: list[date]
    share_mean: np.ndarray   # (T, 6) parts lissées
    share_lo: np.ndarray     # (T, 6) borne 2.5 %
    share_hi: np.ndarray     # (T, 6) borne 97.5 %
    x: np.ndarray            # état filtré à as_of (5,)
    P: np.ndarray            # covariance filtrée (5,5)
    q: float                 # variance quotidienne retenue
    loglik: float
    n_polls: int


def prepare_polls(polls: list[dict], profiles: dict[str, PollsterProfile],
                  parties: list[str] | None = None) -> list[PreparedPoll]:
    ps = parties or settings.parties
    min_main = max(3, len(ps) - 2)
    out = []
    for poll in polls:
        prof = profile_for(profiles, poll["pollster"])
        shares = {p: v for p, v in poll["shares"].items() if p in ps}
        main = {p: v for p, v in shares.items() if p != "AUT"}
        if len(main) < min_main:
            continue
        adj = adjust_shares(shares, prof)
        p_adj = shares_vector(adj, ps)
        n_eff = effective_n(poll.get("sample_size"), poll.get("mode", "unknown"))
        # la cote du sondeur module la précision effective (jamais > ±40 %)
        n_scaled = n_eff * prof.weight_multiplier
        # maison sans historique : erreur excédentaire accrue; les observations
        # spéciales (ex. partielles) fixent leur propre erreur excédentaire
        excess = poll.get("excess_pp") or (
            settings.excess_poll_sd_pp * (1.0 if prof.mae_pp is not None else 1.25))
        R = alr_obs_cov(p_adj, n_scaled, excess_pp=excess)
        out.append(PreparedPoll(
            field_end=poll["field_end"], pollster=poll["pollster"], y=alr(p_adj),
            R=R, p_adj=p_adj, n_eff=n_eff, multiplier=prof.weight_multiplier))
    out.sort(key=lambda p: p.field_end)
    return out


def _filter(prepared: list[PreparedPoll], grid: list[date], q: float,
            x0: np.ndarray, P0: np.ndarray, smooth: bool = False):
    """Filtre de Kalman sur grille quotidienne; lisseur RTS optionnel."""
    T = len(grid)
    Mx = len(x0)
    by_day: dict[date, list[PreparedPoll]] = {}
    for p in prepared:
        by_day.setdefault(p.field_end, []).append(p)

    x, P = x0.copy(), P0.copy()
    loglik = 0.0
    xs_pred = np.zeros((T, Mx)); Ps_pred = np.zeros((T, Mx, Mx))
    xs_filt = np.zeros((T, Mx)); Ps_filt = np.zeros((T, Mx, Mx))
    I = np.eye(Mx)
    for t, d in enumerate(grid):
        if t > 0:
            P = P + q * I
        xs_pred[t], Ps_pred[t] = x, P
        for obs in by_day.get(d, []):
            S = P + obs.R
            Sinv = np.linalg.inv(S)
            innov = obs.y - x
            K = P @ Sinv
            x = x + K @ innov
            P = (I - K) @ P
            P = 0.5 * (P + P.T)
            sign, logdet = np.linalg.slogdet(S)
            loglik += -0.5 * (Mx * np.log(2 * np.pi) + logdet + innov @ Sinv @ innov)
        xs_filt[t], Ps_filt[t] = x, P

    if not smooth:
        return x, P, loglik, None, None
    # Lisseur de Rauch–Tung–Striebel
    xs_s = xs_filt.copy(); Ps_s = Ps_filt.copy()
    for t in range(T - 2, -1, -1):
        G = Ps_filt[t] @ np.linalg.inv(Ps_pred[t + 1])
        xs_s[t] = xs_filt[t] + G @ (xs_s[t + 1] - xs_pred[t + 1])
        Ps_s[t] = Ps_filt[t] + G @ (Ps_s[t + 1] - Ps_pred[t + 1]) @ G.T
    return x, P, loglik, xs_s, Ps_s


def fit_trend(polls: list[dict], profiles: dict[str, PollsterProfile],
              as_of: date, rng: np.random.Generator | None = None,
              parties: list[str] | None = None) -> TrendResult | None:
    ps = parties or settings.parties
    Mx = len(ps) - 1
    prepared = [p for p in prepare_polls(polls, profiles, ps) if p.field_end <= as_of]
    if len(prepared) < 3:
        return None
    rng = rng or np.random.default_rng(20261005)
    start = prepared[0].field_end
    grid = [start + timedelta(days=i) for i in range((as_of - start).days + 1)]

    x0 = prepared[0].y.copy()
    P0 = np.eye(Mx) * 0.09  # a priori large (~±6 pp)

    # Estimation de q par vraisemblance profilée (grille log-uniforme)
    qs = np.geomspace(settings.rw_daily_sd_min ** 2, settings.rw_daily_sd_max ** 2, 12)
    lls = []
    for q in qs:
        _, _, ll, _, _ = _filter(prepared, grid, float(q), x0, P0)
        lls.append(ll)
    q_best = float(qs[int(np.argmax(lls))])

    x, P, ll, xs_s, Ps_s = _filter(prepared, grid, q_best, x0, P0, smooth=True)

    # Parts lissées + intervalle crédible 95 % par échantillonnage
    T = len(grid)
    share_mean = np.zeros((T, Mx + 1)); share_lo = np.zeros((T, Mx + 1)); share_hi = np.zeros((T, Mx + 1))
    n_draw = 400
    for t in range(T):
        L = np.linalg.cholesky(Ps_s[t] + 1e-10 * np.eye(Mx))
        draws = xs_s[t] + rng.standard_normal((n_draw, Mx)) @ L.T
        s = inv_alr(draws)
        share_mean[t] = inv_alr(xs_s[t])
        share_lo[t] = np.percentile(s, 2.5, axis=0)
        share_hi[t] = np.percentile(s, 97.5, axis=0)

    return TrendResult(dates=grid, share_mean=share_mean, share_lo=share_lo,
                       share_hi=share_hi, x=x, P=P, q=q_best, loglik=float(ll),
                       n_polls=len(prepared))