//! Extended Kalman Filter implementation.
//!
//! Best for systems where you have a good mathematical model.
//!
//! A "good" mathematical model is mono-modal, and could be described
//! in closed-form calculus expressions. Neither of these are strict
//! requirements, but the Extended Kalman Filter works best when
//! they're true or nearly so. The Particle Filter is a better choice
//! if these restrictions are too confining.

use super::*;
use nalgebra::{Const, DefaultAllocator, DimName, OMatrix, OVector, allocator::Allocator};

/// Describes how your system changes over time.
///
/// Implement this trait by filling in `predict_next` and `process_noise_std`.
/// The filter figures out the rest automatically.
///
/// # Type Parameters
/// - `N`: Number of state variables (e.g., 3 for [x, y, z])
/// - `C`: Number of control inputs (e.g., 2 for [thrust, steering])
///
/// # Example
///
/// ```
/// use estimators::ekf::ProcessModel;
///
/// struct BoatModel;
///
/// impl ProcessModel<3, 2> for BoatModel {
///     fn predict_next(&self, state: &[f64; 3], control: &[f64; 2], dt: f64) -> [f64; 3] {
///         let [x, y, heading] = *state;
///         let [throttle, rudder] = *control;
///         let speed = throttle * 5.0;  // max 5 m/s
///
///         [
///             x + speed * heading.cos() * dt,
///             y + speed * heading.sin() * dt,
///             heading + rudder * dt,
///         ]
///     }
///
///     fn process_noise_std(&self) -> [f64; 3] {
///         [0.1, 0.1, 0.05]  // [x noise, y noise, heading noise]
///     }
/// }
///
/// // Test the model
/// let model = BoatModel;
/// let state = [0.0, 0.0, 0.0];  // At origin, heading east
/// let control = [1.0, 0.0];     // Full throttle, no rudder
/// let next = model.predict_next(&state, &control, 1.0);
/// assert!((next[0] - 5.0).abs() < 0.01);  // Moved 5m east
/// ```
pub trait ProcessModel<const N: usize, const C: usize> {
    /// Predict the next state given current state, control input, and time step.
    ///
    /// This is the heart of your model. Just describe what happens to your
    /// system over a small time step.
    fn predict_next(&self, state: &[f64; N], control: &[f64; C], dt: f64) -> [f64; N];

    /// How much random noise affects each state variable (per second).
    ///
    /// Return the standard deviation for each variable. For example, if your
    /// position drifts by about ±0.1 meters per second due to wind or other
    /// random factors, use 0.1.
    fn process_noise_std(&self) -> [f64; N];

    /// Advanced: Override this if you want to compute the sensitivity matrix
    /// analytically instead of using automatic numerical differentiation.
    ///
    /// Most users can ignore this — the default implementation calculates
    /// this automatically by testing how small changes in each state variable
    /// affect the prediction.
    fn prediction_sensitivity(
        &self,
        state: &[f64; N],
        control: &[f64; C],
        dt: f64,
    ) -> [[f64; N]; N] {
        numerical_jacobian(|s| self.predict_next(s, control, dt), state)
    }
}

/// Describes what your sensors measure.
///
/// Implement this trait by filling in `expected_measurement` and `measurement_noise_std`.
///
/// # Type Parameters
/// - `N`: Number of state variables
/// - `M`: Number of sensor readings
///
/// # Example
///
/// ```
/// use estimators::ekf::MeasurementModel;
///
/// struct GpsSensor;
///
/// // State is [x, y, heading], but GPS only measures [x, y]
/// impl MeasurementModel<3, 2> for GpsSensor {
///     fn expected_measurement(&self, state: &[f64; 3]) -> [f64; 2] {
///         [state[0], state[1]]  // GPS reads x and y position
///     }
///
///     fn measurement_noise_std(&self) -> [f64; 2] {
///         [2.5, 2.5]  // GPS accurate to about ±2.5 meters
///     }
/// }
///
/// // Test the sensor model
/// let sensor = GpsSensor;
/// let state = [10.0, 20.0, 1.57];  // At (10, 20), heading north
/// let expected = sensor.expected_measurement(&state);
/// assert_eq!(expected, [10.0, 20.0]);
/// ```
pub trait MeasurementModel<const N: usize, const M: usize> {
    /// What measurement would you expect given a particular state?
    ///
    /// If you knew the true state exactly, what would your sensor report?
    fn expected_measurement(&self, state: &[f64; N]) -> [f64; M];

    /// How noisy is each sensor reading?
    ///
    /// Return the standard deviation for each measurement. For example,
    /// if your GPS is accurate to about ±2.5 meters, use 2.5.
    fn measurement_noise_std(&self) -> [f64; M];

    /// Advanced: Override this if you want to compute the sensitivity matrix
    /// analytically instead of using automatic numerical differentiation.
    ///
    /// Most users can ignore this — the default implementation calculates
    /// this automatically by testing how small changes in each state variable
    /// affect the expected measurement.
    fn measurement_sensitivity(&self, state: &[f64; N]) -> [[f64; M]; N] {
        numerical_jacobian(|s| self.expected_measurement(s), state)
    }
}

/// Compute a Jacobian matrix using numerical differentiation.
///
/// Uses central differences for accuracy: for each input variable,
/// we nudge it slightly up and down and measure how the output changes.
pub fn numerical_jacobian<const N: usize, const M: usize, F>(
    f: F,
    state: &[f64; N],
) -> [[f64; M]; N]
where
    F: Fn(&[f64; N]) -> [f64; M],
{
    const EPSILON: f64 = 1e-8;
    let mut jacobian = [[0.0; M]; N];

    for i in 0..N {
        let mut state_plus = *state;
        let mut state_minus = *state;
        state_plus[i] += EPSILON;
        state_minus[i] -= EPSILON;

        let f_plus = f(&state_plus);
        let f_minus = f(&state_minus);

        for j in 0..M {
            jacobian[i][j] = (f_plus[j] - f_minus[j]) / (2.0 * EPSILON);
        }
    }

    jacobian
}

/// Convert array-of-arrays to nalgebra matrix (transposed to match convention).
fn jacobian_to_matrix<const ROWS: usize, const COLS: usize>(
    jac: [[f64; ROWS]; COLS],
) -> OMatrix<f64, Const<ROWS>, Const<COLS>>
where
    Const<ROWS>: DimName,
    Const<COLS>: DimName,
    DefaultAllocator: Allocator<Const<ROWS>, Const<COLS>>,
{
    let mut mat = OMatrix::<f64, Const<ROWS>, Const<COLS>>::zeros();
    for col in 0..COLS {
        for row in 0..ROWS {
            mat[(row, col)] = jac[col][row];
        }
    }
    mat
}

/// Build a diagonal covariance matrix from standard deviations.
fn std_to_covariance<const N: usize>(std_devs: [f64; N]) -> OMatrix<f64, Const<N>, Const<N>>
where
    Const<N>: DimName,
    DefaultAllocator: Allocator<Const<N>, Const<N>>,
{
    let mut cov = OMatrix::<f64, Const<N>, Const<N>>::zeros();
    for i in 0..N {
        cov[(i, i)] = std_devs[i] * std_devs[i];
    }
    cov
}

/// Extended Kalman Filter.
///
/// Tracks state by combining predictions from your model with noisy measurements.
///
/// # Example
///
/// ```
/// use estimators::ekf::{ProcessModel, MeasurementModel, ExtendedKalmanFilter};
/// use estimators::StateEstimator;
///
/// // Simple constant-velocity model: state is [position, velocity]
/// struct ConstantVelocity;
///
/// impl ProcessModel<2, 0> for ConstantVelocity {
///     fn predict_next(&self, state: &[f64; 2], _control: &[f64; 0], dt: f64) -> [f64; 2] {
///         [state[0] + state[1] * dt, state[1]]
///     }
///     fn process_noise_std(&self) -> [f64; 2] {
///         [0.1, 0.01]
///     }
/// }
///
/// struct PositionSensor;
///
/// impl MeasurementModel<2, 1> for PositionSensor {
///     fn expected_measurement(&self, state: &[f64; 2]) -> [f64; 1] {
///         [state[0]]  // We only measure position
///     }
///     fn measurement_noise_std(&self) -> [f64; 1] {
///         [0.5]
///     }
/// }
///
/// // Create the filter
/// let mut ekf = ExtendedKalmanFilter::new(
///     [0.0, 1.0],    // Start at position 0, velocity 1
///     [1.0, 0.5],    // Initial uncertainty
///     ConstantVelocity,
///     PositionSensor,
/// );
///
/// // Run a few iterations
/// for i in 0..5 {
///     ekf.predict(&[], 0.1);  // 100ms time step, no control input
///     let measurement = [(i as f64 + 1.0) * 0.1];  // Simulated measurement
///     ekf.update(&measurement).unwrap();
/// }
///
/// let state = ekf.estimate();
/// println!("Final estimate: pos={:.2}, vel={:.2}", state[0], state[1]);
/// ```
pub struct ExtendedKalmanFilter<const N: usize, const M: usize, const C: usize, PM, MM>
where
    Const<N>: DimName,
    Const<M>: DimName,
    Const<C>: DimName,
    DefaultAllocator: Allocator<Const<N>>
        + Allocator<Const<N>, Const<N>>
        + Allocator<Const<M>>
        + Allocator<Const<M>, Const<N>>
        + Allocator<Const<N>, Const<M>>
        + Allocator<Const<M>, Const<M>>,
{
    state: OVector<f64, Const<N>>,
    covariance: OMatrix<f64, Const<N>, Const<N>>,
    process_model: PM,
    measurement_model: MM,
}

impl<const N: usize, const M: usize, const C: usize, PM, MM> ExtendedKalmanFilter<N, M, C, PM, MM>
where
    Const<N>: DimName,
    Const<M>: DimName,
    Const<C>: DimName,
    PM: ProcessModel<N, C>,
    MM: MeasurementModel<N, M>,
    DefaultAllocator: Allocator<Const<N>>
        + Allocator<Const<N>, Const<N>>
        + Allocator<Const<M>>
        + Allocator<Const<M>, Const<N>>
        + Allocator<Const<N>, Const<M>>
        + Allocator<Const<M>, Const<M>>,
{
    /// Create a new Extended Kalman Filter.
    ///
    /// # Arguments
    /// * `initial_state` - Your best guess at the starting state
    /// * `initial_uncertainty` - How uncertain you are about each state variable
    ///   (standard deviation). Use larger values if you're very unsure.
    /// * `process_model` - Describes how your system evolves
    /// * `measurement_model` - Describes what your sensors measure
    pub fn new(
        initial_state: [f64; N],
        initial_uncertainty: [f64; N],
        process_model: PM,
        measurement_model: MM,
    ) -> Self {
        Self {
            state: OVector::<f64, Const<N>>::from_column_slice(&initial_state),
            covariance: std_to_covariance(initial_uncertainty),
            process_model,
            measurement_model,
        }
    }

    /// Create with a custom initial covariance matrix (advanced).
    pub fn with_covariance(
        initial_state: [f64; N],
        initial_covariance: [[f64; N]; N],
        process_model: PM,
        measurement_model: MM,
    ) -> Self {
        let mut cov = OMatrix::<f64, Const<N>, Const<N>>::zeros();
        for i in 0..N {
            for j in 0..N {
                cov[(i, j)] = initial_covariance[i][j];
            }
        }

        Self {
            state: OVector::<f64, Const<N>>::from_column_slice(&initial_state),
            covariance: cov,
            process_model,
            measurement_model,
        }
    }

    /// Reset the filter to a known state.
    ///
    /// Useful if you have external information (like a GPS fix) that you
    /// want to use as a fresh starting point.
    pub fn reset(&mut self, state: [f64; N], uncertainty: [f64; N]) {
        self.state = OVector::<f64, Const<N>>::from_column_slice(&state);
        self.covariance = std_to_covariance(uncertainty);
    }
}

impl<const N: usize, const M: usize, const C: usize, PM, MM> StateEstimator
    for ExtendedKalmanFilter<N, M, C, PM, MM>
where
    Const<N>: DimName,
    Const<M>: DimName,
    Const<C>: DimName,
    PM: ProcessModel<N, C>,
    MM: MeasurementModel<N, M>,
    DefaultAllocator: Allocator<Const<N>>
        + Allocator<Const<N>, Const<N>>
        + Allocator<Const<M>>
        + Allocator<Const<M>, Const<N>>
        + Allocator<Const<N>, Const<M>>
        + Allocator<Const<M>, Const<M>>,
{
    type State = [f64; N];
    type Measurement = [f64; M];
    type Control = [f64; C];

    fn predict(&mut self, control: &[f64; C], dt: f64) {
        // Get current state as array
        let state_arr: [f64; N] = self.state.as_slice().try_into().unwrap();

        // Predict next state
        let next_state = self.process_model.predict_next(&state_arr, control, dt);
        self.state = OVector::<f64, Const<N>>::from_column_slice(&next_state);

        // Get sensitivity matrix and convert to nalgebra
        let sensitivity = self
            .process_model
            .prediction_sensitivity(&state_arr, control, dt);
        let f = jacobian_to_matrix::<N, N>(sensitivity);

        // Build process noise covariance (scaled by dt)
        let noise_std = self.process_model.process_noise_std();
        let mut q = OMatrix::<f64, Const<N>, Const<N>>::zeros();
        for i in 0..N {
            q[(i, i)] = noise_std[i] * noise_std[i] * dt;
        }

        // Update covariance: P = F * P * F^T + Q
        self.covariance = &f * &self.covariance * f.transpose() + q;
    }

    fn update(&mut self, measurement: &[f64; M]) -> Result<()> {
        let state_arr: [f64; N] = self.state.as_slice().try_into().unwrap();

        // What measurement would we expect?
        let expected = self.measurement_model.expected_measurement(&state_arr);
        let z_pred = OVector::<f64, Const<M>>::from_column_slice(&expected);
        let z = OVector::<f64, Const<M>>::from_column_slice(measurement);

        // Innovation (difference between actual and expected)
        let innovation = z - z_pred;

        // Measurement sensitivity matrix
        let sensitivity = self.measurement_model.measurement_sensitivity(&state_arr);
        let h = jacobian_to_matrix::<M, N>(sensitivity);

        // Measurement noise covariance
        let r = std_to_covariance(self.measurement_model.measurement_noise_std());

        // Innovation covariance: S = H * P * H^T + R
        let s = &h * &self.covariance * h.transpose() + &r;

        // Kalman gain: K = P * H^T * S^(-1)
        let s_inv = s
            .try_inverse()
            .ok_or(EstimatorError::SingularInnovationCovariance)?;
        let k = &self.covariance * h.transpose() * s_inv;

        // Update state
        self.state = &self.state + &k * innovation;

        // Update covariance (Joseph form for numerical stability)
        let i = OMatrix::<f64, Const<N>, Const<N>>::identity();
        let i_kh = &i - &k * &h;
        self.covariance = &i_kh * &self.covariance * i_kh.transpose() + &k * r * k.transpose();

        // Check for numerical instability
        for i in 0..N {
            if !self.state[i].is_finite() {
                return Err(EstimatorError::NumericalInstability(format!(
                    "state[{}] = {}",
                    i, self.state[i]
                )));
            }
        }

        Ok(())
    }

    fn estimate(&self) -> [f64; N] {
        self.state.as_slice().try_into().unwrap()
    }
}

impl<const N: usize, const M: usize, const C: usize, PM, MM> UncertaintyEstimator
    for ExtendedKalmanFilter<N, M, C, PM, MM>
where
    Const<N>: DimName,
    Const<M>: DimName,
    Const<C>: DimName,
    PM: ProcessModel<N, C>,
    MM: MeasurementModel<N, M>,
    DefaultAllocator: Allocator<Const<N>>
        + Allocator<Const<N>, Const<N>>
        + Allocator<Const<M>>
        + Allocator<Const<M>, Const<N>>
        + Allocator<Const<N>, Const<M>>
        + Allocator<Const<M>, Const<M>>,
{
    fn uncertainty(&self) -> Vec<f64> {
        (0..N).map(|i| self.covariance[(i, i)].sqrt()).collect()
    }
}