//! Kalman Filter Suite for state estimation and signal extraction.
//!
//! Implements:
//! - **KalmanFilter**: Standard linear Kalman filter with configurable
//!   transition (F), observation (H), process noise (Q), and measurement
//!   noise (R) matrices. Supports predict/update cycle and log-likelihood.
//! - **UnscentedKF**: Unscented Kalman Filter (sigma-point variant) for
//!   nonlinear state estimation using the scaled unscented transform.
//!
//! All linear algebra is implemented from scratch in Rust -- no external
//! crate dependencies beyond PyO3.
//!
//! Pro tier required.

use pyo3::prelude::*;
use std::sync::Mutex;

// ===========================================================================
// Private linear algebra helpers
// ===========================================================================

/// Matrix type: row-major Vec<Vec<f64>>.
type Mat = Vec<Vec<f64>>;

/// Create an n x m zero matrix.
#[inline]
fn mat_zeros(n: usize, m: usize) -> Mat {
    vec![vec![0.0; m]; n]
}

/// Create an n x n identity matrix.
#[inline]
fn mat_eye(n: usize) -> Mat {
    let mut m = mat_zeros(n, n);
    for i in 0..n {
        m[i][i] = 1.0;
    }
    m
}

/// Matrix multiplication: A (n x k) * B (k x m) -> C (n x m).
fn mat_mul(a: &Mat, b: &Mat) -> Mat {
    let n = a.len();
    if n == 0 {
        return vec![];
    }
    let k = a[0].len();
    let m = b[0].len();
    debug_assert!(b.len() == k, "mat_mul: dimension mismatch");
    let mut c = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            let mut sum = 0.0;
            for l in 0..k {
                sum += a[i][l] * b[l][j];
            }
            c[i][j] = sum;
        }
    }
    c
}

/// Matrix-vector multiplication: A (n x m) * v (m) -> w (n).
fn mat_vec_mul(a: &Mat, v: &[f64]) -> Vec<f64> {
    let n = a.len();
    let mut w = vec![0.0; n];
    for i in 0..n {
        let mut sum = 0.0;
        for j in 0..a[i].len() {
            sum += a[i][j] * v[j];
        }
        w[i] = sum;
    }
    w
}

/// Matrix addition: A + B (element-wise, same dimensions).
fn mat_add(a: &Mat, b: &Mat) -> Mat {
    let n = a.len();
    let m = a[0].len();
    let mut c = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            c[i][j] = a[i][j] + b[i][j];
        }
    }
    c
}

/// Matrix subtraction: A - B (element-wise, same dimensions).
fn mat_sub(a: &Mat, b: &Mat) -> Mat {
    let n = a.len();
    let m = a[0].len();
    let mut c = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            c[i][j] = a[i][j] - b[i][j];
        }
    }
    c
}

/// Matrix transpose: A^T.
fn mat_transpose(a: &Mat) -> Mat {
    let n = a.len();
    if n == 0 {
        return vec![];
    }
    let m = a[0].len();
    let mut t = mat_zeros(m, n);
    for i in 0..n {
        for j in 0..m {
            t[j][i] = a[i][j];
        }
    }
    t
}

/// Scalar multiplication: c * A.
fn mat_scale(c: f64, a: &Mat) -> Mat {
    let n = a.len();
    if n == 0 {
        return vec![];
    }
    let m = a[0].len();
    let mut r = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            r[i][j] = c * a[i][j];
        }
    }
    r
}

/// Matrix inverse via Gauss-Jordan elimination with partial pivoting.
/// Returns None if the matrix is singular (det ~ 0).
fn mat_inv(a: &Mat) -> Option<Mat> {
    let n = a.len();
    if n == 0 {
        return Some(vec![]);
    }
    // Build augmented matrix [A | I]
    let mut aug = mat_zeros(n, 2 * n);
    for i in 0..n {
        for j in 0..n {
            aug[i][j] = a[i][j];
        }
        aug[i][n + i] = 1.0;
    }

    for col in 0..n {
        // Partial pivoting: find the row with the largest absolute value in this column
        let mut max_val = aug[col][col].abs();
        let mut max_row = col;
        for row in (col + 1)..n {
            let val = aug[row][col].abs();
            if val > max_val {
                max_val = val;
                max_row = row;
            }
        }
        if max_val < 1e-15 {
            return None; // Singular
        }
        if max_row != col {
            aug.swap(col, max_row);
        }

        // Scale pivot row
        let pivot = aug[col][col];
        for j in 0..(2 * n) {
            aug[col][j] /= pivot;
        }

        // Eliminate column in all other rows
        for row in 0..n {
            if row == col {
                continue;
            }
            let factor = aug[row][col];
            for j in 0..(2 * n) {
                aug[row][j] -= factor * aug[col][j];
            }
        }
    }

    // Extract the inverse from the right half of the augmented matrix
    let mut inv = mat_zeros(n, n);
    for i in 0..n {
        for j in 0..n {
            inv[i][j] = aug[i][n + j];
        }
    }
    Some(inv)
}

/// Cholesky decomposition: A = L * L^T.
/// Returns lower triangular matrix L, or None if A is not positive definite.
fn cholesky(a: &Mat) -> Option<Mat> {
    let n = a.len();
    let mut l = mat_zeros(n, n);
    for i in 0..n {
        for j in 0..=i {
            let mut sum = 0.0;
            if j == i {
                for k in 0..j {
                    sum += l[j][k] * l[j][k];
                }
                let diag = a[j][j] - sum;
                if diag < 0.0 {
                    // Not positive definite -- apply small regularization and retry
                    return None;
                }
                l[i][j] = diag.sqrt();
            } else {
                for k in 0..j {
                    sum += l[i][k] * l[j][k];
                }
                if l[j][j].abs() < 1e-15 {
                    return None;
                }
                l[i][j] = (a[i][j] - sum) / l[j][j];
            }
        }
    }
    Some(l)
}

/// Cholesky with regularization fallback: if initial decomposition fails,
/// add increasing diagonal jitter until it succeeds.
fn cholesky_regularized(a: &Mat) -> Mat {
    if let Some(l) = cholesky(a) {
        return l;
    }
    let n = a.len();
    let mut reg = a.clone();
    let mut jitter = 1e-10;
    for _ in 0..10 {
        for i in 0..n {
            reg[i][i] = a[i][i] + jitter;
        }
        if let Some(l) = cholesky(&reg) {
            return l;
        }
        jitter *= 10.0;
    }
    // Last resort: return identity scaled by sqrt of average diagonal
    let avg_diag: f64 = (0..n).map(|i| a[i][i].abs()).sum::<f64>() / n as f64;
    let s = avg_diag.sqrt().max(1e-6);
    let mut l = mat_zeros(n, n);
    for i in 0..n {
        l[i][i] = s;
    }
    l
}

/// Outer product of two vectors: v * w^T.
fn outer_product(v: &[f64], w: &[f64]) -> Mat {
    let n = v.len();
    let m = w.len();
    let mut o = mat_zeros(n, m);
    for i in 0..n {
        for j in 0..m {
            o[i][j] = v[i] * w[j];
        }
    }
    o
}

/// Vector subtraction.
fn vec_sub(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter().zip(b.iter()).map(|(x, y)| x - y).collect()
}

/// Vector addition.
fn vec_add(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()
}

// ===========================================================================
// Standard Kalman Filter
// ===========================================================================

struct KalmanInner {
    state_dim: usize,
    obs_dim: usize,
    /// State vector (state_dim x 1)
    x: Vec<f64>,
    /// State covariance (state_dim x state_dim)
    p: Mat,
    /// State transition matrix F (state_dim x state_dim)
    f: Mat,
    /// Observation matrix H (obs_dim x state_dim)
    h: Mat,
    /// Process noise covariance Q (state_dim x state_dim)
    q: Mat,
    /// Measurement noise covariance R (obs_dim x obs_dim)
    r: Mat,
}

impl KalmanInner {
    fn new(state_dim: usize, obs_dim: usize) -> Self {
        Self {
            state_dim,
            obs_dim,
            x: vec![0.0; state_dim],
            p: mat_eye(state_dim),
            f: mat_eye(state_dim),
            h: {
                // Default H: identity-like mapping (obs_dim rows from state_dim)
                let mut h = mat_zeros(obs_dim, state_dim);
                for i in 0..obs_dim.min(state_dim) {
                    h[i][i] = 1.0;
                }
                h
            },
            q: mat_scale(0.01, &mat_eye(state_dim)),
            r: mat_scale(1.0, &mat_eye(obs_dim)),
        }
    }

    /// Predict step: x_pred = F * x, P_pred = F * P * F^T + Q
    fn predict(&mut self) -> Vec<f64> {
        // x = F * x
        self.x = mat_vec_mul(&self.f, &self.x);
        // P = F * P * F^T + Q
        let fp = mat_mul(&self.f, &self.p);
        let ft = mat_transpose(&self.f);
        self.p = mat_add(&mat_mul(&fp, &ft), &self.q);
        self.x.clone()
    }

    /// Update step with measurement z.
    /// Returns the updated state estimate.
    fn update(&mut self, z: &[f64]) -> Vec<f64> {
        // Innovation: y = z - H * x
        let hx = mat_vec_mul(&self.h, &self.x);
        let y: Vec<f64> = z.iter().zip(hx.iter()).map(|(a, b)| a - b).collect();

        // Innovation covariance: S = H * P * H^T + R
        let hp = mat_mul(&self.h, &self.p);
        let ht = mat_transpose(&self.h);
        let s = mat_add(&mat_mul(&hp, &ht), &self.r);

        // Kalman gain: K = P * H^T * S^{-1}
        let s_inv = mat_inv(&s).unwrap_or_else(|| {
            // Fallback: use R^{-1} scaled (degenerate case)
            mat_scale(1e-6, &mat_eye(self.obs_dim))
        });
        let pht = mat_mul(&self.p, &ht);
        let k = mat_mul(&pht, &s_inv);

        // State update: x = x + K * y
        let ky = mat_vec_mul(&k, &y);
        self.x = vec_add(&self.x, &ky);

        // Covariance update: P = (I - K * H) * P  (Joseph form for numerical stability)
        let kh = mat_mul(&k, &self.h);
        let i_kh = mat_sub(&mat_eye(self.state_dim), &kh);
        let i_kh_p = mat_mul(&i_kh, &self.p);
        let i_kh_t = mat_transpose(&i_kh);
        let kr = mat_mul(&k, &self.r);
        let kt = mat_transpose(&k);
        let krkt = mat_mul(&kr, &kt);
        self.p = mat_add(&mat_mul(&i_kh_p, &i_kh_t), &krkt);

        self.x.clone()
    }

    /// Log-likelihood of a measurement under the current predicted distribution.
    /// Must be called between predict() and update().
    fn log_likelihood(&self, z: &[f64]) -> f64 {
        let hx = mat_vec_mul(&self.h, &self.x);
        let y: Vec<f64> = z.iter().zip(hx.iter()).map(|(a, b)| a - b).collect();

        // S = H * P * H^T + R
        let hp = mat_mul(&self.h, &self.p);
        let ht = mat_transpose(&self.h);
        let s = mat_add(&mat_mul(&hp, &ht), &self.r);

        let k = self.obs_dim as f64;
        let log_2pi = (2.0 * std::f64::consts::PI).ln();

        // log|S| via diagonal of Cholesky
        let l = cholesky_regularized(&s);
        let log_det: f64 = 2.0 * (0..self.obs_dim).map(|i| l[i][i].abs().max(1e-300).ln()).sum::<f64>();

        // y^T * S^{-1} * y
        let s_inv = mat_inv(&s).unwrap_or_else(|| mat_scale(1e-6, &mat_eye(self.obs_dim)));
        let s_inv_y = mat_vec_mul(&s_inv, &y);
        let quad: f64 = y.iter().zip(s_inv_y.iter()).map(|(a, b)| a * b).sum();

        -0.5 * (k * log_2pi + log_det + quad)
    }
}

// ---------------------------------------------------------------------------
// PyO3 interface: KalmanFilter
// ---------------------------------------------------------------------------

/// Standard linear Kalman filter with configurable matrices.
///
/// State estimation for linear dynamical systems:
///   x_{t+1} = F * x_t + w,  w ~ N(0, Q)
///   z_t     = H * x_t + v,  v ~ N(0, R)
///
/// Pro tier required.
#[pyclass]
pub struct KalmanFilter {
    inner: Mutex<KalmanInner>,
}

#[pymethods]
impl KalmanFilter {
    /// Create a new Kalman filter.
    ///
    /// Args:
    ///     state_dim: Dimension of the state vector.
    ///     obs_dim: Dimension of the observation vector.
    #[new]
    #[pyo3(signature = (state_dim, obs_dim))]
    fn new(state_dim: usize, obs_dim: usize) -> PyResult<Self> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if state_dim == 0 || obs_dim == 0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "state_dim and obs_dim must be >= 1",
            ));
        }
        if state_dim > 64 || obs_dim > 64 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "state_dim and obs_dim must be <= 64",
            ));
        }
        Ok(Self {
            inner: Mutex::new(KalmanInner::new(state_dim, obs_dim)),
        })
    }

    /// Set the state transition matrix F (state_dim x state_dim).
    fn set_transition(&self, f_matrix: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if f_matrix.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "F must have {} rows, got {}",
                inner.state_dim,
                f_matrix.len()
            )));
        }
        for (i, row) in f_matrix.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "F row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "F contains non-finite values",
                    ));
                }
            }
        }
        inner.f = f_matrix;
        Ok(())
    }

    /// Set the observation matrix H (obs_dim x state_dim).
    fn set_observation(&self, h_matrix: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if h_matrix.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "H must have {} rows, got {}",
                inner.obs_dim,
                h_matrix.len()
            )));
        }
        for (i, row) in h_matrix.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "H row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "H contains non-finite values",
                    ));
                }
            }
        }
        inner.h = h_matrix;
        Ok(())
    }

    /// Set the process noise covariance Q (state_dim x state_dim).
    fn set_process_noise(&self, q: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if q.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Q must have {} rows, got {}",
                inner.state_dim,
                q.len()
            )));
        }
        for (i, row) in q.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Q row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "Q contains non-finite values",
                    ));
                }
            }
        }
        inner.q = q;
        Ok(())
    }

    /// Set the measurement noise covariance R (obs_dim x obs_dim).
    fn set_measurement_noise(&self, r: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if r.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "R must have {} rows, got {}",
                inner.obs_dim,
                r.len()
            )));
        }
        for (i, row) in r.iter().enumerate() {
            if row.len() != inner.obs_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "R row {} must have {} cols, got {}",
                    i,
                    inner.obs_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "R contains non-finite values",
                    ));
                }
            }
        }
        inner.r = r;
        Ok(())
    }

    /// Set the initial state vector.
    fn set_state(&self, state: Vec<f64>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if state.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "state must have {} elements, got {}",
                inner.state_dim,
                state.len()
            )));
        }
        for &v in &state {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "state contains non-finite values",
                ));
            }
        }
        inner.x = state;
        Ok(())
    }

    /// Set the initial covariance matrix.
    fn set_covariance(&self, cov: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if cov.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "covariance must have {} rows, got {}",
                inner.state_dim,
                cov.len()
            )));
        }
        for (i, row) in cov.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "covariance row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
        }
        inner.p = cov;
        Ok(())
    }

    /// Predict step: propagate state and covariance through transition model.
    ///
    /// Returns:
    ///     Predicted state vector.
    fn predict(&self) -> PyResult<Vec<f64>> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        Ok(inner.predict())
    }

    /// Update step: incorporate a new measurement.
    ///
    /// Args:
    ///     measurement: Observation vector (obs_dim).
    ///
    /// Returns:
    ///     Updated state vector.
    fn update(&self, measurement: Vec<f64>) -> PyResult<Vec<f64>> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if measurement.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "measurement must have {} elements, got {}",
                inner.obs_dim,
                measurement.len()
            )));
        }
        for &v in &measurement {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "measurement contains non-finite values",
                ));
            }
        }
        Ok(inner.update(&measurement))
    }

    /// Get the current state vector.
    fn state(&self) -> Vec<f64> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.x.clone()
    }

    /// Get the current covariance matrix.
    fn covariance(&self) -> Vec<Vec<f64>> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.p.clone()
    }

    /// Compute the log-likelihood of a measurement under the current
    /// predicted distribution. Call between predict() and update().
    ///
    /// Args:
    ///     measurement: Observation vector (obs_dim).
    ///
    /// Returns:
    ///     Log-likelihood value.
    fn log_likelihood(&self, measurement: Vec<f64>) -> PyResult<f64> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if measurement.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "measurement must have {} elements, got {}",
                inner.obs_dim,
                measurement.len()
            )));
        }
        Ok(inner.log_likelihood(&measurement))
    }

    fn __repr__(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        format!(
            "KalmanFilter(state_dim={}, obs_dim={})",
            inner.state_dim, inner.obs_dim
        )
    }
}

// ===========================================================================
// Unscented Kalman Filter
// ===========================================================================

struct UKFInner {
    state_dim: usize,
    obs_dim: usize,
    /// State vector
    x: Vec<f64>,
    /// State covariance
    p: Mat,
    /// State transition matrix F (linearized -- applied as matrix multiply)
    f: Mat,
    /// Observation matrix H
    h: Mat,
    /// Process noise covariance Q
    q: Mat,
    /// Measurement noise covariance R
    r: Mat,
    /// Sigma point scaling parameters
    alpha: f64,
    beta: f64,
    kappa: f64,
    /// Precomputed weights
    wm: Vec<f64>, // mean weights
    wc: Vec<f64>, // covariance weights
    lambda: f64,
}

impl UKFInner {
    fn new(state_dim: usize, obs_dim: usize, alpha: f64, beta: f64, kappa: f64) -> Self {
        let n = state_dim as f64;
        let lambda = alpha * alpha * (n + kappa) - n;

        // 2*n + 1 sigma points
        let n_sigma = 2 * state_dim + 1;
        let mut wm = vec![0.0; n_sigma];
        let mut wc = vec![0.0; n_sigma];

        // Weight for the mean sigma point
        wm[0] = lambda / (n + lambda);
        wc[0] = lambda / (n + lambda) + (1.0 - alpha * alpha + beta);

        // Weights for the remaining sigma points
        let w = 1.0 / (2.0 * (n + lambda));
        for i in 1..n_sigma {
            wm[i] = w;
            wc[i] = w;
        }

        Self {
            state_dim,
            obs_dim,
            x: vec![0.0; state_dim],
            p: mat_eye(state_dim),
            f: mat_eye(state_dim),
            h: {
                let mut h = mat_zeros(obs_dim, state_dim);
                for i in 0..obs_dim.min(state_dim) {
                    h[i][i] = 1.0;
                }
                h
            },
            q: mat_scale(0.01, &mat_eye(state_dim)),
            r: mat_scale(1.0, &mat_eye(obs_dim)),
            alpha,
            beta,
            kappa,
            wm,
            wc,
            lambda,
        }
    }

    /// Generate sigma points from current state and covariance.
    /// Returns a Vec of (2*n+1) sigma point vectors.
    fn generate_sigma_points(&self) -> Vec<Vec<f64>> {
        let n = self.state_dim;
        let n_sigma = 2 * n + 1;
        let mut sigma_points = Vec::with_capacity(n_sigma);

        // Compute sqrt((n + lambda) * P) via Cholesky
        let scale = (n as f64 + self.lambda).max(0.0);
        let scaled_p = mat_scale(scale, &self.p);
        let l = cholesky_regularized(&scaled_p);

        // Sigma point 0: the mean
        sigma_points.push(self.x.clone());

        // Sigma points 1..n: x + L_i (i-th column of L)
        for i in 0..n {
            let col: Vec<f64> = (0..n).map(|r| l[r][i]).collect();
            sigma_points.push(vec_add(&self.x, &col));
        }

        // Sigma points n+1..2n: x - L_i
        for i in 0..n {
            let col: Vec<f64> = (0..n).map(|r| l[r][i]).collect();
            sigma_points.push(vec_sub(&self.x, &col));
        }

        sigma_points
    }

    /// Propagate sigma points through the state transition model (F * x).
    fn propagate_state(&self, sigma_points: &[Vec<f64>]) -> Vec<Vec<f64>> {
        sigma_points
            .iter()
            .map(|sp| mat_vec_mul(&self.f, sp))
            .collect()
    }

    /// Propagate sigma points through the observation model (H * x).
    fn propagate_observation(&self, sigma_points: &[Vec<f64>]) -> Vec<Vec<f64>> {
        sigma_points
            .iter()
            .map(|sp| mat_vec_mul(&self.h, sp))
            .collect()
    }

    /// Compute weighted mean of a set of sigma points.
    fn weighted_mean(&self, points: &[Vec<f64>]) -> Vec<f64> {
        let dim = points[0].len();
        let mut mean = vec![0.0; dim];
        for (i, pt) in points.iter().enumerate() {
            for j in 0..dim {
                mean[j] += self.wm[i] * pt[j];
            }
        }
        mean
    }

    /// Compute weighted covariance: Sum wc_i * (p_i - mean) * (p_i - mean)^T + noise.
    fn weighted_covariance(&self, points: &[Vec<f64>], mean: &[f64], noise: &Mat) -> Mat {
        let dim = mean.len();
        let mut cov = mat_zeros(dim, dim);
        for (i, pt) in points.iter().enumerate() {
            let diff = vec_sub(pt, mean);
            let op = outer_product(&diff, &diff);
            for r in 0..dim {
                for c in 0..dim {
                    cov[r][c] += self.wc[i] * op[r][c];
                }
            }
        }
        mat_add(&cov, noise)
    }

    /// Compute weighted cross-covariance: Sum wc_i * (x_i - x_mean) * (z_i - z_mean)^T.
    fn weighted_cross_covariance(
        &self,
        x_points: &[Vec<f64>],
        x_mean: &[f64],
        z_points: &[Vec<f64>],
        z_mean: &[f64],
    ) -> Mat {
        let xdim = x_mean.len();
        let zdim = z_mean.len();
        let mut cc = mat_zeros(xdim, zdim);
        for (i, (xp, zp)) in x_points.iter().zip(z_points.iter()).enumerate() {
            let xdiff = vec_sub(xp, x_mean);
            let zdiff = vec_sub(zp, z_mean);
            for r in 0..xdim {
                for c in 0..zdim {
                    cc[r][c] += self.wc[i] * xdiff[r] * zdiff[c];
                }
            }
        }
        cc
    }

    /// Predict step.
    fn predict(&mut self) -> Vec<f64> {
        // Generate sigma points from current state
        let sigma_pts = self.generate_sigma_points();

        // Propagate through state model
        let pred_pts = self.propagate_state(&sigma_pts);

        // Compute predicted mean and covariance
        self.x = self.weighted_mean(&pred_pts);
        self.p = self.weighted_covariance(&pred_pts, &self.x, &self.q);

        self.x.clone()
    }

    /// Update step with measurement.
    fn update(&mut self, z: &[f64]) -> Vec<f64> {
        // Generate sigma points from predicted state
        let sigma_pts = self.generate_sigma_points();

        // Propagate through observation model
        let obs_pts = self.propagate_observation(&sigma_pts);

        // Predicted measurement mean
        let z_pred = self.weighted_mean(&obs_pts);

        // Innovation covariance: S = Pzz + R
        let s = self.weighted_covariance(&obs_pts, &z_pred, &self.r);

        // Cross-covariance: Pxz
        let pxz = self.weighted_cross_covariance(&sigma_pts, &self.x, &obs_pts, &z_pred);

        // Kalman gain: K = Pxz * S^{-1}
        let s_inv = mat_inv(&s).unwrap_or_else(|| mat_scale(1e-6, &mat_eye(self.obs_dim)));
        let k = mat_mul(&pxz, &s_inv);

        // Innovation
        let innovation = vec_sub(z, &z_pred);

        // State update
        let k_inn = mat_vec_mul(&k, &innovation);
        self.x = vec_add(&self.x, &k_inn);

        // Covariance update: P = P - K * S * K^T
        let ks = mat_mul(&k, &s);
        let kt = mat_transpose(&k);
        let kskt = mat_mul(&ks, &kt);
        self.p = mat_sub(&self.p, &kskt);

        self.x.clone()
    }

    /// Log-likelihood of a measurement under the predicted distribution.
    fn log_likelihood(&self, z: &[f64]) -> f64 {
        // Generate sigma points
        let sigma_pts = self.generate_sigma_points();
        let obs_pts = self.propagate_observation(&sigma_pts);
        let z_pred = self.weighted_mean(&obs_pts);
        let s = self.weighted_covariance(&obs_pts, &z_pred, &self.r);

        let innovation = vec_sub(z, &z_pred);
        let k = self.obs_dim as f64;
        let log_2pi = (2.0 * std::f64::consts::PI).ln();

        let l = cholesky_regularized(&s);
        let log_det: f64 =
            2.0 * (0..self.obs_dim).map(|i| l[i][i].abs().max(1e-300).ln()).sum::<f64>();

        let s_inv = mat_inv(&s).unwrap_or_else(|| mat_scale(1e-6, &mat_eye(self.obs_dim)));
        let s_inv_y = mat_vec_mul(&s_inv, &innovation);
        let quad: f64 = innovation.iter().zip(s_inv_y.iter()).map(|(a, b)| a * b).sum();

        -0.5 * (k * log_2pi + log_det + quad)
    }
}

// ---------------------------------------------------------------------------
// PyO3 interface: UnscentedKF
// ---------------------------------------------------------------------------

/// Unscented Kalman Filter using the scaled unscented transform.
///
/// Handles nonlinear dynamics better than the standard KF by using
/// deterministic sigma points instead of linearization.
///
/// Pro tier required.
#[pyclass]
pub struct UnscentedKF {
    inner: Mutex<UKFInner>,
}

#[pymethods]
impl UnscentedKF {
    /// Create a new Unscented Kalman Filter.
    ///
    /// Args:
    ///     state_dim: Dimension of the state vector.
    ///     obs_dim: Dimension of the observation vector.
    ///     alpha: Spread of sigma points (default 0.001). Controls how far
    ///            sigma points are from the mean. Smaller = tighter.
    ///     beta: Prior distribution parameter (default 2.0). For Gaussian, 2 is optimal.
    ///     kappa: Secondary scaling parameter (default 0.0). Typically 0 or 3 - n.
    #[new]
    #[pyo3(signature = (state_dim, obs_dim, alpha=0.001, beta=2.0, kappa=0.0))]
    fn new(state_dim: usize, obs_dim: usize, alpha: f64, beta: f64, kappa: f64) -> PyResult<Self> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if state_dim == 0 || obs_dim == 0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "state_dim and obs_dim must be >= 1",
            ));
        }
        if state_dim > 64 || obs_dim > 64 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "state_dim and obs_dim must be <= 64",
            ));
        }
        if alpha <= 0.0 || !alpha.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "alpha must be > 0 and finite",
            ));
        }
        Ok(Self {
            inner: Mutex::new(UKFInner::new(state_dim, obs_dim, alpha, beta, kappa)),
        })
    }

    /// Set the state transition matrix F (state_dim x state_dim).
    fn set_transition(&self, f_matrix: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if f_matrix.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "F must have {} rows, got {}",
                inner.state_dim,
                f_matrix.len()
            )));
        }
        for (i, row) in f_matrix.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "F row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
        }
        inner.f = f_matrix;
        Ok(())
    }

    /// Set the observation matrix H (obs_dim x state_dim).
    fn set_observation(&self, h_matrix: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if h_matrix.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "H must have {} rows, got {}",
                inner.obs_dim,
                h_matrix.len()
            )));
        }
        for (i, row) in h_matrix.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "H row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
        }
        inner.h = h_matrix;
        Ok(())
    }

    /// Set the process noise covariance Q (state_dim x state_dim).
    fn set_process_noise(&self, q: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if q.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Q must have {} rows, got {}",
                inner.state_dim,
                q.len()
            )));
        }
        for (i, row) in q.iter().enumerate() {
            if row.len() != inner.state_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Q row {} must have {} cols, got {}",
                    i,
                    inner.state_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "Q contains non-finite values",
                    ));
                }
            }
        }
        inner.q = q;
        Ok(())
    }

    /// Set the measurement noise covariance R (obs_dim x obs_dim).
    fn set_measurement_noise(&self, r: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if r.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "R must have {} rows, got {}",
                inner.obs_dim,
                r.len()
            )));
        }
        for (i, row) in r.iter().enumerate() {
            if row.len() != inner.obs_dim {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "R row {} must have {} cols, got {}",
                    i,
                    inner.obs_dim,
                    row.len()
                )));
            }
            for &v in row {
                if !v.is_finite() {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "R contains non-finite values",
                    ));
                }
            }
        }
        inner.r = r;
        Ok(())
    }

    /// Set the initial state vector.
    fn set_state(&self, state: Vec<f64>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if state.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "state must have {} elements, got {}",
                inner.state_dim,
                state.len()
            )));
        }
        for &v in &state {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "state contains non-finite values",
                ));
            }
        }
        inner.x = state;
        Ok(())
    }

    /// Set the initial covariance matrix.
    fn set_covariance(&self, cov: Vec<Vec<f64>>) -> PyResult<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if cov.len() != inner.state_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "covariance must have {} rows, got {}",
                inner.state_dim,
                cov.len()
            )));
        }
        inner.p = cov;
        Ok(())
    }

    /// Predict step: propagate sigma points through state model.
    ///
    /// Returns:
    ///     Predicted state vector.
    fn predict(&self) -> PyResult<Vec<f64>> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        Ok(inner.predict())
    }

    /// Update step: incorporate a new measurement via sigma points.
    ///
    /// Args:
    ///     measurement: Observation vector (obs_dim).
    ///
    /// Returns:
    ///     Updated state vector.
    fn update(&self, measurement: Vec<f64>) -> PyResult<Vec<f64>> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if measurement.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "measurement must have {} elements, got {}",
                inner.obs_dim,
                measurement.len()
            )));
        }
        for &v in &measurement {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "measurement contains non-finite values",
                ));
            }
        }
        Ok(inner.update(&measurement))
    }

    /// Get the current state vector.
    fn state(&self) -> Vec<f64> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.x.clone()
    }

    /// Get the current covariance matrix.
    fn covariance(&self) -> Vec<Vec<f64>> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.p.clone()
    }

    /// Compute the log-likelihood of a measurement under the current
    /// predicted distribution.
    ///
    /// Args:
    ///     measurement: Observation vector (obs_dim).
    ///
    /// Returns:
    ///     Log-likelihood value.
    fn log_likelihood(&self, measurement: Vec<f64>) -> PyResult<f64> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if measurement.len() != inner.obs_dim {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "measurement must have {} elements, got {}",
                inner.obs_dim,
                measurement.len()
            )));
        }
        Ok(inner.log_likelihood(&measurement))
    }

    fn __repr__(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        format!(
            "UnscentedKF(state_dim={}, obs_dim={}, alpha={:.4}, beta={:.1}, kappa={:.1})",
            inner.state_dim, inner.obs_dim, inner.alpha, inner.beta, inner.kappa
        )
    }
}

// ===========================================================================
// Registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<KalmanFilter>()?;
    m.add_class::<UnscentedKF>()?;
    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

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

    // -----------------------------------------------------------------------
    // Linear algebra helpers
    // -----------------------------------------------------------------------

    #[test]
    fn test_mat_eye() {
        let eye = mat_eye(3);
        assert_eq!(eye[0][0], 1.0);
        assert_eq!(eye[1][1], 1.0);
        assert_eq!(eye[2][2], 1.0);
        assert_eq!(eye[0][1], 0.0);
    }

    #[test]
    fn test_mat_mul() {
        let a = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
        let b = vec![vec![5.0, 6.0], vec![7.0, 8.0]];
        let c = mat_mul(&a, &b);
        assert_eq!(c[0][0], 19.0);
        assert_eq!(c[0][1], 22.0);
        assert_eq!(c[1][0], 43.0);
        assert_eq!(c[1][1], 50.0);
    }

    #[test]
    fn test_mat_transpose() {
        let a = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
        let at = mat_transpose(&a);
        assert_eq!(at.len(), 3);
        assert_eq!(at[0].len(), 2);
        assert_eq!(at[0][0], 1.0);
        assert_eq!(at[1][0], 2.0);
        assert_eq!(at[2][1], 6.0);
    }

    #[test]
    fn test_mat_inv_identity() {
        let eye = mat_eye(3);
        let inv = mat_inv(&eye).unwrap();
        for i in 0..3 {
            for j in 0..3 {
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!((inv[i][j] - expected).abs() < 1e-10);
            }
        }
    }

    #[test]
    fn test_mat_inv_2x2() {
        // [[4, 7], [2, 6]] -> inv = [[ 6, -7], [-2,  4]] / 10
        let a = vec![vec![4.0, 7.0], vec![2.0, 6.0]];
        let inv = mat_inv(&a).unwrap();
        let det = 4.0 * 6.0 - 7.0 * 2.0; // 10
        assert!((inv[0][0] - 6.0 / det).abs() < 1e-10);
        assert!((inv[0][1] - (-7.0 / det)).abs() < 1e-10);
        assert!((inv[1][0] - (-2.0 / det)).abs() < 1e-10);
        assert!((inv[1][1] - 4.0 / det).abs() < 1e-10);
    }

    #[test]
    fn test_mat_inv_singular() {
        let a = vec![vec![1.0, 2.0], vec![2.0, 4.0]];
        assert!(mat_inv(&a).is_none());
    }

    #[test]
    fn test_cholesky() {
        // Positive definite: [[4, 2], [2, 3]]
        let a = vec![vec![4.0, 2.0], vec![2.0, 3.0]];
        let l = cholesky(&a).unwrap();
        // L * L^T should equal A
        let lt = mat_transpose(&l);
        let result = mat_mul(&l, &lt);
        for i in 0..2 {
            for j in 0..2 {
                assert!(
                    (result[i][j] - a[i][j]).abs() < 1e-10,
                    "cholesky mismatch at [{i}][{j}]"
                );
            }
        }
    }

    #[test]
    fn test_cholesky_not_pd() {
        // Not positive definite
        let a = vec![vec![1.0, 2.0], vec![2.0, 1.0]];
        assert!(cholesky(&a).is_none());
    }

    // -----------------------------------------------------------------------
    // Standard Kalman Filter
    // -----------------------------------------------------------------------

    #[test]
    fn test_kalman_1d_constant() {
        // 1D constant model: state = position, observation = position + noise
        let mut kf = KalmanInner::new(1, 1);
        kf.f = vec![vec![1.0]]; // state transition: position stays
        kf.h = vec![vec![1.0]]; // we observe position directly
        kf.q = vec![vec![0.001]]; // small process noise
        kf.r = vec![vec![1.0]]; // measurement noise

        // Feed noisy observations of value 5.0
        let observations = vec![5.2, 4.8, 5.1, 4.9, 5.0, 5.3, 4.7, 5.1, 5.0, 4.9];
        for &z in &observations {
            kf.predict();
            kf.update(&[z]);
        }

        // After 10 observations, estimate should be close to 5.0
        assert!((kf.x[0] - 5.0).abs() < 0.5, "estimate = {}", kf.x[0]);
        // Covariance should have decreased from initial 1.0
        assert!(kf.p[0][0] < 1.0, "cov = {}", kf.p[0][0]);
    }

    #[test]
    fn test_kalman_2d_velocity() {
        // 2D state: [position, velocity]
        // Observation: position only
        let mut kf = KalmanInner::new(2, 1);
        let dt = 1.0;
        kf.f = vec![vec![1.0, dt], vec![0.0, 1.0]]; // constant velocity model
        kf.h = vec![vec![1.0, 0.0]]; // observe position only
        kf.q = vec![
            vec![0.25 * dt.powi(4), 0.5 * dt.powi(3)],
            vec![0.5 * dt.powi(3), dt * dt],
        ];
        kf.q = mat_scale(0.01, &kf.q); // small process noise
        kf.r = vec![vec![0.5]]; // measurement noise

        // True trajectory: position = 10 + 2*t (velocity = 2)
        for t in 0..20 {
            let true_pos = 10.0 + 2.0 * t as f64;
            let noise = (t as f64 * 0.37).sin() * 0.5; // deterministic "noise"
            let z = true_pos + noise;
            kf.predict();
            kf.update(&[z]);
        }

        // Position should be near 10 + 2*19 = 48
        assert!(
            (kf.x[0] - 48.0).abs() < 3.0,
            "position = {}",
            kf.x[0]
        );
        // Velocity should be near 2.0
        assert!(
            (kf.x[1] - 2.0).abs() < 1.0,
            "velocity = {}",
            kf.x[1]
        );
    }

    #[test]
    fn test_kalman_log_likelihood() {
        let mut kf = KalmanInner::new(1, 1);
        kf.r = vec![vec![1.0]];
        kf.predict();

        // Observation at the predicted mean should have higher likelihood
        // than one far away
        let ll_at_mean = kf.log_likelihood(&[0.0]);
        let ll_far = kf.log_likelihood(&[10.0]);
        assert!(ll_at_mean > ll_far);
    }

    // -----------------------------------------------------------------------
    // Unscented Kalman Filter
    // -----------------------------------------------------------------------

    #[test]
    fn test_ukf_sigma_points() {
        let ukf = UKFInner::new(2, 1, 0.001, 2.0, 0.0);
        let pts = ukf.generate_sigma_points();
        // 2*n + 1 = 5 sigma points
        assert_eq!(pts.len(), 5);
        // First point should be the mean
        assert_eq!(pts[0], ukf.x);
    }

    #[test]
    fn test_ukf_weights_sum() {
        let ukf = UKFInner::new(3, 1, 0.5, 2.0, 0.0);
        let wm_sum: f64 = ukf.wm.iter().sum();
        let wc_sum: f64 = ukf.wc.iter().sum();
        // Mean weights should sum to 1
        assert!(
            (wm_sum - 1.0).abs() < 1e-10,
            "wm sum = {}",
            wm_sum
        );
        // Covariance weights sum depends on beta
        assert!(wc_sum.is_finite());
    }

    #[test]
    fn test_ukf_1d_constant() {
        let mut ukf = UKFInner::new(1, 1, 0.1, 2.0, 0.0);
        ukf.q = vec![vec![0.001]];
        ukf.r = vec![vec![1.0]];

        let observations = vec![5.2, 4.8, 5.1, 4.9, 5.0, 5.3, 4.7, 5.1, 5.0, 4.9];
        for &z in &observations {
            ukf.predict();
            ukf.update(&[z]);
        }

        assert!((ukf.x[0] - 5.0).abs() < 0.5, "ukf estimate = {}", ukf.x[0]);
    }

    #[test]
    fn test_ukf_2d_velocity() {
        let mut ukf = UKFInner::new(2, 1, 0.1, 2.0, 0.0);
        let dt = 1.0;
        ukf.f = vec![vec![1.0, dt], vec![0.0, 1.0]];
        ukf.h = vec![vec![1.0, 0.0]];
        ukf.q = mat_scale(0.01, &vec![
            vec![0.25 * dt.powi(4), 0.5 * dt.powi(3)],
            vec![0.5 * dt.powi(3), dt * dt],
        ]);
        ukf.r = vec![vec![0.5]];

        for t in 0..20 {
            let true_pos = 10.0 + 2.0 * t as f64;
            let noise = (t as f64 * 0.37).sin() * 0.5;
            let z = true_pos + noise;
            ukf.predict();
            ukf.update(&[z]);
        }

        assert!(
            (ukf.x[0] - 48.0).abs() < 3.0,
            "ukf position = {}",
            ukf.x[0]
        );
        assert!(
            (ukf.x[1] - 2.0).abs() < 1.0,
            "ukf velocity = {}",
            ukf.x[1]
        );
    }

    #[test]
    fn test_ukf_log_likelihood() {
        let mut ukf = UKFInner::new(1, 1, 0.1, 2.0, 0.0);
        ukf.r = vec![vec![1.0]];
        ukf.predict();

        let ll_at_mean = ukf.log_likelihood(&[0.0]);
        let ll_far = ukf.log_likelihood(&[10.0]);
        assert!(ll_at_mean > ll_far);
    }

    #[test]
    fn test_outer_product() {
        let v = vec![1.0, 2.0];
        let w = vec![3.0, 4.0, 5.0];
        let o = outer_product(&v, &w);
        assert_eq!(o.len(), 2);
        assert_eq!(o[0].len(), 3);
        assert_eq!(o[0][0], 3.0);
        assert_eq!(o[1][2], 10.0);
    }
}
