oxictl 0.1.1

use crate::core::matrix::{matmul, matvec, outer, Matrix};
use crate::core::scalar::ControlScalar;

/// Ensemble Kalman Filter (EnKF).
///
/// A Monte Carlo implementation of the Kalman Filter for nonlinear systems.
/// Instead of propagating a covariance matrix analytically, an **ensemble** of
/// `E` state vectors is propagated through the (possibly nonlinear) dynamics.
/// The sample covariance of the ensemble approximates `P`.
///
/// **Perturbed-observation update** (Burgers et al., 1998):
/// Each ensemble member receives a different, randomly perturbed observation
/// `z_e = z + ε_e` (where `ε_e ~ N(0, R)`) so that the ensemble covariance
/// contracts correctly upon assimilation.
///
/// Deterministic perturbation generation uses a **Linear Congruential Generator
/// (LCG)** seeded from user-supplied state — no `rand` crate dependency.
///
/// # Type Parameters
/// * `S` – scalar type (`f32` or `f64`)
/// * `N` – state dimension
/// * `M` – measurement dimension
/// * `E` – ensemble size (must be ≥ 2)
#[derive(Debug, Clone)]
pub struct EnsembleKf<S: ControlScalar, const N: usize, const M: usize, const E: usize> {
    /// Process noise covariance (N×N) — used for ensemble perturbation at init.
    pub q: Matrix<S, N, N>,
    /// Measurement noise covariance (M×M).
    pub r: Matrix<S, M, M>,
    /// Nonlinear state-transition function `f(x, u) -> x_next`.
    f: fn(&[S; N], &[S; 1]) -> [S; N],
    /// Measurement function `h(x) -> z`.
    h: fn(&[S; N]) -> [S; M],
    /// Ensemble members — each column is one member (N × E stored row-major).
    ensemble: [[S; N]; E],
    /// LCG state for deterministic pseudo-random number generation.
    lcg_state: u64,
}

/// Error type for the Ensemble Kalman Filter.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EnkfError {
    /// Ensemble size is too small (need E ≥ 2).
    EnsembleTooSmall,
    /// Innovation covariance matrix is singular.
    SingularInnovationCovariance,
    /// Process noise Cholesky failed (Q not positive definite).
    QNotPositiveDefinite,
}

impl core::fmt::Display for EnkfError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            EnkfError::EnsembleTooSmall => write!(f, "EnKF: ensemble size must be ≥ 2"),
            EnkfError::SingularInnovationCovariance => {
                write!(f, "EnKF: innovation covariance is singular")
            }
            EnkfError::QNotPositiveDefinite => write!(f, "EnKF: Q not positive definite"),
        }
    }
}

/// Linear Congruential Generator — produces pseudo-random u64 values.
///
/// Uses Knuth's parameters: `a = 6364136223846793005`, `c = 1442695040888963407`.
#[inline]
fn lcg_next(state: &mut u64) -> u64 {
    *state = state
        .wrapping_mul(6_364_136_223_846_793_005)
        .wrapping_add(1_442_695_040_888_963_407);
    *state
}

/// Convert a u64 LCG sample to a standard-normal variate via Box-Muller.
///
/// Two LCG calls are consumed per two normal samples.  Returns `(n1, n2)`.
fn lcg_normal_pair<S: ControlScalar>(state: &mut u64) -> (S, S) {
    // Map two integers to (0,1) uniforms
    let u1_raw = lcg_next(state);
    let u2_raw = lcg_next(state);

    // Map to (0, 1) — avoid exactly 0
    let u1 = S::from_f64((u1_raw >> 11) as f64 / (1u64 << 53) as f64).max(S::from_f64(1e-15_f64));
    let u2 = S::from_f64((u2_raw >> 11) as f64 / (1u64 << 53) as f64);

    // Box-Muller transform
    let mag = (-S::TWO * u1.ln()).sqrt();
    let angle = S::TWO * S::PI * u2;
    (mag * angle.cos(), mag * angle.sin())
}

impl<S: ControlScalar, const N: usize, const M: usize, const E: usize> EnsembleKf<S, N, M, E> {
    /// Create an Ensemble Kalman Filter.
    ///
    /// The initial ensemble is generated by perturbing `x0` with samples from
    /// `N(0, P0)` using a deterministic LCG seeded by `seed`.
    ///
    /// Returns `Err` if `E < 2` or if `P0` is not positive definite.
    pub fn new(
        q: Matrix<S, N, N>,
        r: Matrix<S, M, M>,
        f: fn(&[S; N], &[S; 1]) -> [S; N],
        h: fn(&[S; N]) -> [S; M],
        x0: [S; N],
        p0: Matrix<S, N, N>,
        seed: u64,
    ) -> Result<Self, EnkfError> {
        if E < 2 {
            return Err(EnkfError::EnsembleTooSmall);
        }

        let l0 = p0.cholesky().ok_or(EnkfError::QNotPositiveDefinite)?;
        let mut lcg_state = seed ^ 0xDEAD_BEEF_CAFE_1234;

        // Generate ensemble by perturbing x0
        let ensemble: [[S; N]; E] = core::array::from_fn(|_| {
            // Draw N normal samples and form perturbation = L0 · ε
            let mut eps = [S::ZERO; N];
            let mut idx = 0;
            while idx < N {
                let (n1, n2) = lcg_normal_pair::<S>(&mut lcg_state);
                eps[idx] = n1;
                if idx + 1 < N {
                    eps[idx + 1] = n2;
                }
                idx += 2;
            }
            let perturbation = matvec(&l0, &eps);
            core::array::from_fn(|i| x0[i] + perturbation[i])
        });

        Ok(Self {
            q,
            r,
            f,
            h,
            ensemble,
            lcg_state,
        })
    }

    /// **Predict step** — propagate each ensemble member through `f`.
    ///
    /// Process noise is added by sampling from `N(0, Q)` using the internal LCG.
    ///
    /// Returns `Err` if `Q` is not positive definite.
    pub fn predict(&mut self, u: S) -> Result<(), EnkfError> {
        let lq = self.q.cholesky().ok_or(EnkfError::QNotPositiveDefinite)?;
        let u_arr = [u];

        for member in self.ensemble.iter_mut() {
            // Sample process noise: w ~ N(0, Q)
            let mut eps = [S::ZERO; N];
            let mut idx = 0;
            while idx < N {
                let (n1, n2) = lcg_normal_pair::<S>(&mut self.lcg_state);
                eps[idx] = n1;
                if idx + 1 < N {
                    eps[idx + 1] = n2;
                }
                idx += 2;
            }
            let w = matvec(&lq, &eps);

            let x_next = (self.f)(member, &u_arr);
            for i in 0..N {
                member[i] = x_next[i] + w[i];
            }
        }
        Ok(())
    }

    /// **Update step** — perturbed-observation EnKF update.
    ///
    /// 1. Compute ensemble mean `x̄` and anomaly matrix.
    /// 2. Compute sample covariance `P_e` and cross-covariance `P_xz`.
    /// 3. Perturb observation: `z_e = z + ε`, `ε ~ N(0, R)`.
    /// 4. Update each member: `x_e ← x_e + K · (z_e - H·x_e)`.
    ///
    /// Returns the innovation `z - H·x̄`, or `Err` if innovation covariance
    /// is singular.
    pub fn update(&mut self, z: &[S; M]) -> Result<[S; M], EnkfError> {
        // ---- Ensemble mean ----
        let mut x_mean = [S::ZERO; N];
        for member in &self.ensemble {
            for i in 0..N {
                x_mean[i] += member[i];
            }
        }
        let inv_e = S::ONE / S::from_f64(E as f64);
        for xi in x_mean.iter_mut() {
            *xi *= inv_e;
        }

        // ---- Ensemble measurement mean ----
        let mut z_mean = [S::ZERO; M];
        for member in &self.ensemble {
            let z_e = (self.h)(member);
            for j in 0..M {
                z_mean[j] += z_e[j];
            }
        }
        for zj in z_mean.iter_mut() {
            *zj *= inv_e;
        }

        // ---- Sample covariances P_zz and P_xz ----
        // P_zz = (1/(E-1)) Σ (H·x_e - z_mean)(H·x_e - z_mean)^T  +  R
        // P_xz = (1/(E-1)) Σ (x_e - x_mean)(H·x_e - z_mean)^T
        let inv_e1 = S::ONE / S::from_f64((E - 1) as f64);
        let mut p_zz = self.r; // start with R
        let mut p_xz = Matrix::<S, N, M>::zeros();

        for member in &self.ensemble {
            let z_e = (self.h)(member);
            let dz: [S; M] = core::array::from_fn(|j| z_e[j] - z_mean[j]);
            let dx: [S; N] = core::array::from_fn(|i| member[i] - x_mean[i]);

            let dz_dzt = outer(&dz, &dz).scale(inv_e1);
            let dx_dzt = outer(&dx, &dz).scale(inv_e1);
            p_zz = p_zz.add_mat(&dz_dzt);
            p_xz = p_xz.add_mat(&dx_dzt);
        }

        // Kalman gain: K = P_xz · P_zz⁻¹
        let p_zz_inv = p_zz.inv().ok_or(EnkfError::SingularInnovationCovariance)?;
        let k = matmul(&p_xz, &p_zz_inv);

        // Cholesky of R for perturbation generation
        let lr = self.r.cholesky();

        // ---- Update each ensemble member ----
        for member in self.ensemble.iter_mut() {
            // Perturb observation: z_pert = z + ε, ε ~ N(0, R)
            let z_pert: [S; M] = if let Some(ref lr_mat) = lr {
                let mut eps = [S::ZERO; M];
                let mut idx = 0;
                while idx < M {
                    let (n1, n2) = lcg_normal_pair::<S>(&mut self.lcg_state);
                    eps[idx] = n1;
                    if idx + 1 < M {
                        eps[idx + 1] = n2;
                    }
                    idx += 2;
                }
                let noise = matvec(lr_mat, &eps);
                core::array::from_fn(|j| z[j] + noise[j])
            } else {
                *z
            };

            // Innovation for this member
            let hx_e = (self.h)(member);
            let innov_e: [S; M] = core::array::from_fn(|j| z_pert[j] - hx_e[j]);

            // x_e ← x_e + K · innov_e
            let k_innov = matvec(&k, &innov_e);
            for i in 0..N {
                member[i] += k_innov[i];
            }
        }

        // Return innovation w.r.t. ensemble mean (diagnostic)
        let hz_mean = (self.h)(&x_mean);
        let innovation: [S; M] = core::array::from_fn(|j| z[j] - hz_mean[j]);
        Ok(innovation)
    }

    /// Ensemble mean (current state estimate).
    pub fn state(&self) -> [S; N] {
        let mut mean = [S::ZERO; N];
        for member in &self.ensemble {
            for i in 0..N {
                mean[i] += member[i];
            }
        }
        let inv_e = S::ONE / S::from_f64(E as f64);
        core::array::from_fn(|i| mean[i] * inv_e)
    }

    /// Sample covariance of the ensemble.
    pub fn covariance(&self) -> Matrix<S, N, N> {
        let mean = self.state();
        let inv_e1 = S::ONE / S::from_f64((E - 1) as f64);
        let mut p = Matrix::<S, N, N>::zeros();
        for member in &self.ensemble {
            let dx: [S; N] = core::array::from_fn(|i| member[i] - mean[i]);
            let dx_dxt = outer(&dx, &dx).scale(inv_e1);
            p = p.add_mat(&dx_dxt);
        }
        p
    }

    /// Access the raw ensemble (each element is one member of size N).
    pub fn ensemble(&self) -> &[[S; N]; E] {
        &self.ensemble
    }

    /// Effective sample size:  `N_eff = (Σ w_i)² / Σ w_i²`.
    ///
    /// For a uniform-weight ensemble (EnKF), this is always `E`.  Provided for
    /// API symmetry with particle filters.
    pub fn effective_sample_size(&self) -> S {
        S::from_f64(E as f64)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Simple linear dynamics: position-velocity
    fn f_pv(x: &[f64; 2], u: &[f64; 1]) -> [f64; 2] {
        [x[0] + 0.01 * x[1] + 0.0 * u[0], x[1]]
    }

    fn h_pos(x: &[f64; 2]) -> [f64; 1] {
        [x[0]]
    }

    fn build_enkf() -> EnsembleKf<f64, 2, 1, 20> {
        let q = Matrix::<f64, 2, 2>::identity().scale(1e-4);
        let r = Matrix::<f64, 1, 1>::identity().scale(0.1);
        let p0 = Matrix::<f64, 2, 2>::identity();
        EnsembleKf::new(q, r, f_pv, h_pos, [0.0_f64; 2], p0, 42).expect("valid EnKF")
    }

    #[test]
    fn new_creates_ensemble() {
        let enkf = build_enkf();
        assert_eq!(enkf.ensemble().len(), 20);
    }

    #[test]
    fn new_fails_for_ensemble_size_1() {
        let q = Matrix::<f64, 2, 2>::identity().scale(1e-4);
        let r = Matrix::<f64, 1, 1>::identity().scale(0.1);
        let p0 = Matrix::<f64, 2, 2>::identity();
        // E = 1 should fail; but E is a const param so we use E=2 min:
        // We test via the EnsembleTooSmall path which is compile-time unreachable
        // for E>=2. Instead verify that E=2 works fine.
        let result: Result<EnsembleKf<f64, 2, 1, 2>, _> =
            EnsembleKf::new(q, r, f_pv, h_pos, [0.0_f64; 2], p0, 1);
        assert!(result.is_ok());
    }

    #[test]
    fn predict_runs() {
        let mut enkf = build_enkf();
        assert!(enkf.predict(0.0).is_ok());
    }

    #[test]
    fn update_returns_innovation() {
        let mut enkf = build_enkf();
        enkf.predict(0.0).expect("predict");
        let innov = enkf.update(&[1.0]).expect("update");
        assert_eq!(innov.len(), 1);
    }

    #[test]
    fn ensemble_mean_close_to_measurement() {
        let mut enkf = build_enkf();
        let true_pos = 3.0_f64;
        for _ in 0..200 {
            enkf.predict(0.0).expect("predict");
            enkf.update(&[true_pos]).expect("update");
        }
        let x = enkf.state();
        assert!(
            (x[0] - true_pos).abs() < 1.0,
            "Expected mean ~{true_pos}, got {}",
            x[0]
        );
    }

    #[test]
    fn covariance_is_symmetric() {
        let mut enkf = build_enkf();
        for _ in 0..20 {
            enkf.predict(0.0).expect("predict");
            enkf.update(&[1.0]).expect("update");
        }
        let p = enkf.covariance();
        for i in 0..2 {
            for j in 0..2 {
                assert!(
                    (p.data[i][j] - p.data[j][i]).abs() < 1e-12,
                    "P not symmetric at ({i},{j})"
                );
            }
        }
    }

    #[test]
    fn effective_sample_size() {
        let enkf = build_enkf();
        assert!((enkf.effective_sample_size() - 20.0_f64).abs() < 1e-10);
    }

    #[test]
    fn different_seeds_give_different_ensembles() {
        let q = Matrix::<f64, 2, 2>::identity().scale(1e-4);
        let r = Matrix::<f64, 1, 1>::identity().scale(0.1);
        let p0 = Matrix::<f64, 2, 2>::identity();

        let e1 =
            EnsembleKf::<f64, 2, 1, 10>::new(q, r, f_pv, h_pos, [0.0_f64; 2], p0, 1).expect("e1");
        let e2 =
            EnsembleKf::<f64, 2, 1, 10>::new(q, r, f_pv, h_pos, [0.0_f64; 2], p0, 99).expect("e2");

        // At least one member must differ
        let any_diff = e1
            .ensemble()
            .iter()
            .zip(e2.ensemble().iter())
            .any(|(m1, m2)| m1[0] != m2[0] || m1[1] != m2[1]);
        assert!(
            any_diff,
            "Different seeds should produce different ensembles"
        );
    }
}