#![allow(non_snake_case)]
use nalgebra::{Matrix2, Matrix5, Vector2, Vector5, Matrix2x5};
use core::f32::consts::PI;
use nalgebra::ComplexField;

/// 5-state EKF for planar vehicle: [px, py, vx, vy, theta]
pub struct Ekf2D {
    pub x: Vector5<f32>,   // state
    pub P: Matrix5<f32>,   // covariance
    pub Q: Matrix5<f32>,   // process noise
    pub R_gps: Matrix2<f32>, // GPS measurement noise (px,py)
}

impl Ekf2D {
    /// Small helper: rotation matrix R(theta)
    fn rotation(theta: f32) -> Matrix2<f32> {
        let c = theta.cos();
        let s = theta.sin();
        Matrix2::new(c, -s, s, c)
    }

    /// Derivative dR/dtheta
    fn rotation_derivative(theta: f32) -> Matrix2<f32> {
        let c = theta.cos();
        let s = theta.sin();
        // derivative of [[c, -s],[s, c]] = [[-s, -c],[c, -s]]
        Matrix2::new(-s, -c, c, -s)
    }

    /// Predict step using IMU readings in body frame:
    /// - accel_b: Vector2 (ax, ay) in body frame (m/s^2)
    /// - gyro_z: yaw rate (rad/s)
    /// - dt: seconds
    pub fn predict(&mut self, accel_b: Vector2<f32>, gyro_z: f32, dt: f32) {
        // Unpack state
        let px = self.x[0];
        let py = self.x[1];
        let vx = self.x[2];
        let vy = self.x[3];
        let theta = self.x[4];

        // rotation and rotated acceleration
        let R = Self::rotation(theta);
        let a_w = R * accel_b; // world-frame acceleration

        // Nonlinear state propagation (discrete)
        let p_new = Vector2::new(px, py) + Vector2::new(vx, vy) * dt + 0.5 * a_w * dt * dt;
        let v_new = Vector2::new(vx, vy) + a_w * dt;
        let theta_new = theta + gyro_z * dt;

        // Updated state
        self.x[0] = p_new[0];
        self.x[1] = p_new[1];
        self.x[2] = v_new[0];
        self.x[3] = v_new[1];
        self.x[4] = theta_new;

        // Build Jacobian F = df/dx (5x5)
        // Using ordering [px, py, vx, vy, theta]
        let mut F = Matrix5::identity();

        // ∂p'/∂v = dt * I2
        F[(0,2)] = dt; F[(1,3)] = dt;

        // ∂p'/∂theta = 0.5 * dt^2 * dR/dtheta * a_b (2x1 into columns)
        let dR_dth = Self::rotation_derivative(theta);
        let dp_dth = 0.5 * dt * dt * (dR_dth * accel_b);
        F[(0,4)] = dp_dth[0];
        F[(1,4)] = dp_dth[1];

        // ∂v'/∂theta = dt * dR/dtheta * a_b
        let dv_dth = dt * (dR_dth * accel_b);
        F[(2,4)] = dv_dth[0];
        F[(3,4)] = dv_dth[1];

        // ∂v'/∂v = I2 (already identity on diagonal)
        // ∂p'/∂p = I2 (already identity)

        // ∂theta'/∂theta = 1 (identity)

        // Covariance propagation: P = F P Fᵀ + Q (we might scale Q by dt if desired)
        self.P = F * self.P * F.transpose() + self.Q;
    }

    /// GPS update: provide measured position (px,py) in world frame
    pub fn update_gps(&mut self, z_px: f32, z_py: f32) {
        // measurement matrix H (2x5) picks px,py
        let mut H = Matrix2x5::zeros();
        H[(0,0)] = 1.0;
        H[(1,1)] = 1.0;

        // y = z - H x
        let z = Vector2::new(z_px, z_py);
        let hx = Vector2::new(self.x[0], self.x[1]);
        let y = z - hx;

        // S = H P Hᵀ + R
        let S = H * self.P * H.transpose() + self.R_gps;

        // K = P Hᵀ S⁻¹  (5x2)
        let K = self.P * H.transpose() * S.try_inverse().unwrap_or_else(|| {
            // fallback to small-regularized inverse if S is singular
            let mut Sreg = S;
            Sreg[(0,0)] += 1e-6;
            Sreg[(1,1)] += 1e-6;
            Sreg.try_inverse().unwrap()
        });

        // x = x + K y
        let dx = K * y;
        self.x += dx;

        // P = (I - K H) P
        let I5 = Matrix5::identity();
        let KH = K * H;
        self.P = (I5 - KH) * self.P;
    }

    /// Optional: set a more informed initial pose when first GPS arrives.
    /// This can reduce large initial uncertainty in position and optionally set heading.
    pub fn initialize_with_gps(&mut self, px: f32, py: f32, pos_variance: f32) {
        self.x[0] = px;
        self.x[1] = py;
        self.P[(0,0)] = pos_variance;
        self.P[(1,1)] = pos_variance;
        // leave velocity and heading as-is (or set to small uncertainty if you know them)
    }

        /// Zero-velocity update (ZUPT). Call this when you know the vehicle is stationary.
    pub fn update_zupt(&mut self) {
        // Measurement matrix H (2x5): extracts vx, vy
        let mut H = Matrix2x5::zeros();
        H[(0,2)] = 1.0;
        H[(1,3)] = 1.0;

        // Expected measurement: current velocity estimate
        let hx = Vector2::new(self.x[2], self.x[3]);

        // "Observed" velocity is zero
        let z = Vector2::zeros();
        let y = z - hx;

        // Choose a small measurement noise R (you can tune this)
        let mut R = Matrix2::zeros();
        R[(0,0)] = 0.01;
        R[(1,1)] = 0.01;

        // S = H P Hᵀ + R
        let S = H * self.P * H.transpose() + R;

        // Kalman gain
        let K = self.P * H.transpose() * S.try_inverse().unwrap();

        // State correction
        self.x += K * y;

        // Covariance correction
        let I5 = Matrix5::identity();
        self.P = (I5 - K * H) * self.P;
    }
}

impl Default for Ekf2D {
    fn default() -> Self {
        let x = Vector5::zeros();
        // default: unknown position/vel/heading. We'll set large covariance
        let mut P = Matrix5::zeros();
        P[(0,0)] = 1e6; // px
        P[(1,1)] = 1e6; // py
        P[(2,2)] = 1e3; // vx
        P[(3,3)] = 1e3; // vy
        P[(4,4)] = (2.0*PI).powi(2); // theta huge uncertainty

        // Process noise - tune to your sensors:
        // Represents uncertainty in the dynamics (accel noise affects v and p, gyro noise affects theta)
        let mut Q = Matrix5::zeros();
        Q[(0,0)] = 0.1;
        Q[(1,1)] = 0.1;
        Q[(2,2)] = 1.0;
        Q[(3,3)] = 1.0;
        Q[(4,4)] = 0.5; // heading diffusion

        let mut R_gps = Matrix2::zeros();
        R_gps[(0,0)] = 5.0; // variance in meters^2 (tune to your GPS)
        R_gps[(1,1)] = 5.0;

        Self { x, P, Q, R_gps }
    }
}