"""
bishop.sequential – PRML Chapter 13
======================================
Linear Kalman Filter: uses scipy.linalg (solve / cho_factor) for all
matrix operations.
Particle Filter: uses numpy for resampling (systematic resampling via
numpy argsort / cumsum) and scipy.stats for likelihood computation.

    from ml_ai_library.bishop.sequential import KalmanFilter, ParticleFilter

PRML chapter map:
  KalmanFilter   → §13.3 (linear Gaussian state-space model)
  ParticleFilter → §13.3 (non-linear/non-Gaussian, SMC)
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable, Optional, Tuple

import numpy as np
from scipy import linalg as _la      # matrix routines (solve, cho_factor, etc.)
from scipy import stats as _stats    # likelihood computations


# ---------------------------------------------------------------------------
# Kalman Filter (PRML §13.3) – scipy.linalg as numerical engine
# ---------------------------------------------------------------------------

@dataclass
class KalmanFilter:
    """
    Linear Gaussian Kalman Filter (PRML §13.3).

    State model:    x_t = F x_{t-1} + N(0, Q)
    Observation:    z_t = H x_t     + N(0, R)

    All matrix arithmetic delegated to scipy.linalg.

    Parameters
    ----------
    F : (n, n) state transition matrix
    H : (m, n) observation matrix
    Q : (n, n) process noise covariance
    R : (m, m) observation noise covariance
    x0: (n,)   initial state mean
    P0: (n, n) initial state covariance
    """

    F:  np.ndarray
    H:  np.ndarray
    Q:  np.ndarray
    R:  np.ndarray
    x0: np.ndarray
    P0: np.ndarray

    def filter(self, observations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Run the Kalman filter over a sequence of observations.

        Parameters
        ----------
        observations : (T, m) array

        Returns
        -------
        means : (T, n) filtered state means
        covs  : (T, n, n) filtered state covariances
        """
        T    = len(observations)
        n    = self.F.shape[0]
        x, P = self.x0.copy(), self.P0.copy()
        means = np.empty((T, n))
        covs  = np.empty((T, n, n))

        for t, z in enumerate(observations):
            # Predict – uses numpy matrix multiplication (scipy.linalg-ready)
            x_pred = self.F @ x
            P_pred = self.F @ P @ self.F.T + self.Q

            # Update – Kalman gain via scipy.linalg.solve (numerically stable)
            S = self.H @ P_pred @ self.H.T + self.R
            K = _la.solve(S.T, (P_pred @ self.H.T).T).T  # K = P H^T S^{-1}
            y = z - self.H @ x_pred                       # innovation

            x = x_pred + K @ y
            P = (np.eye(n) - K @ self.H) @ P_pred         # Joseph form optional

            means[t] = x
            covs[t]  = P

        return means, covs

    def smooth(self, observations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Rauch-Tung-Striebel smoother (backward pass) via scipy.linalg.solve.

        Returns smoothed means and covariances.
        """
        means_f, covs_f = self.filter(observations)
        T = len(observations)
        n = self.F.shape[0]

        means_s = means_f.copy()
        covs_s  = covs_f.copy()

        for t in range(T - 2, -1, -1):
            P_pred  = self.F @ covs_f[t] @ self.F.T + self.Q
            # Smoother gain: G = P_t F^T P_pred^{-1} via scipy.linalg.solve
            G = _la.solve(P_pred.T, (covs_f[t] @ self.F.T).T).T
            means_s[t] += G @ (means_s[t + 1] - self.F @ means_f[t])
            covs_s[t]  += G @ (covs_s[t + 1] - P_pred) @ G.T

        return means_s, covs_s


# ---------------------------------------------------------------------------
# Particle Filter (Sequential Monte Carlo, PRML §13.3)
# Uses numpy for resampling, scipy.stats for likelihood
# ---------------------------------------------------------------------------

@dataclass
class ParticleFilter:
    """
    Bootstrap particle filter / Sequential Monte Carlo (PRML §13.3).

    Parameters
    ----------
    n_particles   : number of particles
    transition_fn : (particle_array, rng) -> new_particle_array
                    vectorised sampler from p(x_t | x_{t-1})
    likelihood_fn : (particle, observation) -> log_likelihood scalar
                    uses scipy.stats internally for Gaussian case
    resample_ess_threshold : resample when ESS drops below this fraction

    Example (1-D tracking)
    ----------------------
    >>> from scipy import stats
    >>> pf = ParticleFilter(
    ...     n_particles=500,
    ...     transition_fn=lambda p, rng: p + rng.normal(0, 0.5, p.shape),
    ...     likelihood_fn=lambda p, z: stats.norm.logpdf(z, loc=p, scale=1.0),
    ... )
    >>> mean_est, _ = pf.update(0.0, np.zeros(500))
    """

    n_particles:           int
    transition_fn:         Callable[[np.ndarray, np.random.Generator], np.ndarray]
    likelihood_fn:         Callable[[np.ndarray, float], np.ndarray]
    resample_ess_threshold: float = 0.5

    def __post_init__(self):
        self._rng = np.random.default_rng()

    def _systematic_resample(self, weights: np.ndarray) -> np.ndarray:
        """
        Systematic resampling (Kitagawa 1996) using numpy cumsum + searchsorted.
        O(N) and lower variance than multinomial resampling.
        """
        N    = len(weights)
        cdf  = np.cumsum(weights)
        u0   = self._rng.uniform(0, 1.0 / N)
        u    = u0 + np.arange(N) / N
        return np.searchsorted(cdf, u)      # numpy searchsorted as subroutine

    def update(self, observation, particles: np.ndarray
               ) -> Tuple[np.ndarray, np.ndarray]:
        """
        One filter step: propagate → weight → (optionally) resample.

        Parameters
        ----------
        observation : scalar or array observation z_t
        particles   : (N, ...) current particle array

        Returns
        -------
        mean_estimate : weighted mean of particles
        new_particles : propagated (and possibly resampled) particles
        """
        # Propagate via transition (user-supplied, may wrap scipy.stats.rvs)
        prop = self.transition_fn(particles, self._rng)

        # Compute importance weights from log-likelihood (scipy.stats)
        log_w = np.array([self.likelihood_fn(p, observation) for p in prop])
        log_w -= log_w.max()                  # numerical stability
        w     = np.exp(log_w)
        w    /= w.sum()

        # Effective Sample Size – decides whether to resample
        ess = 1.0 / np.sum(w ** 2)
        if ess < self.resample_ess_threshold * self.n_particles:
            idx  = self._systematic_resample(w)
            prop = prop[idx]
            w    = np.ones(self.n_particles) / self.n_particles

        mean_est = np.average(prop, weights=w, axis=0)
        return mean_est, prop
