//! SLAM implementation with s a fixed size full (all location a feature) state Kalman estimate.
//!
//! Useful to test Fast SLAM results.
use bayes_estimate::models::{ExtendedLinearObserver, ExtendedLinearPredictor, KalmanState};
use bayes_estimate::noise::CorrelatedNoise;
use nalgebra::{Const, Dyn, OMatrix, SMatrix, SVector};

// A numeric type
type N = f64;

pub struct KalmanSLAM<const DL: usize, const DF: usize, const DZ: usize> {
    /// Location part of state (particle form).
    pub full: KalmanState<N, Dyn>,
}

impl<const DL: usize, const DF: usize, const DZ: usize> KalmanSLAM<DL, DF, DZ> {
    pub fn init_location(&mut self, location: &KalmanState<N, Const<DL>>) {
        self.full.x.fixed_rows_mut::<DL>(0).copy_from(&location.x);
        self.full
            .X
            .fixed_view_mut::<DL, DL>(0, 0)
            .copy_from(&location.X);
    }

    pub fn predict(&mut self, fx: &SMatrix<N, DL, DL>, noise: &CorrelatedNoise<N, Const<DL>>) {
        // extract location part of full
        let mut loc = KalmanState::<N, Const<DL>> {
            x: self.full.x.fixed_rows::<DL>(0).clone_owned(),
            X: self.full.X.fixed_view::<DL, DL>(0, 0).clone_owned(),
        };
        // predict location, independent of map
        let xx = fx * &loc.x;
        ExtendedLinearPredictor::<f64, Const<DL>>::predict(&mut loc, &xx, fx, noise).unwrap();
        // return location to full
        self.full.x.rows_mut(0, DL).copy_from(&loc.x);
        self.full.X.view_mut((0, 0), (DL, DL)).copy_from(&loc.X);
    }

    pub fn observe(
        &mut self,
        feature: usize,
        z: &SVector<N, DZ>,
        hx_loc: &SMatrix<N, DZ, DL>,
        hx_map: &SMatrix<N, DZ, DF>,
        noise: &CorrelatedNoise<N, Const<DZ>>,
    ) {
        let fi = DL + feature * DF;

        // Create an augmented sparse observe model for full states
        let mut hx = OMatrix::zeros_generic(z.shape_generic().0, Dyn(self.full.x.nrows()));
        hx.fixed_view_mut::<DZ, DL>(0, 0).copy_from(&hx_loc);
        hx.fixed_view_mut::<DZ, DF>(0, fi).copy_from(&hx_map);
        self.full
            .observe_innovation(&(z - &hx * &self.full.x), &hx, noise)
            .unwrap();
    }

    pub fn observe_new(
        &mut self,
        feature: usize,
        z: &SVector<N, DZ>,
        h_loc: &SMatrix<N, DF, DL>,
        h_obs: &SMatrix<N, DF, DZ>,
        noise: &CorrelatedNoise<N, Const<DZ>>,
    ) {
        let fi = DL + feature * DF;

        // feature covariance with existing location and features
        // X+ = [0 Ha] X [0 Ha]' + Hb Z Hb'
        let feature_var = h_loc * &self.full.X.fixed_rows::<DL>(0);
        self.full.X.fixed_rows_mut::<DF>(fi).copy_from(&feature_var);
        self.full
            .X
            .fixed_columns_mut::<DF>(fi)
            .tr_copy_from(&feature_var);

        // feature state and covariance
        let feature_state = h_loc * &self.full.x.fixed_rows::<DL>(0) + h_obs * z;
        self.full
            .x
            .fixed_rows_mut::<DF>(fi)
            .copy_from(&feature_state);
        let feature_cov = h_loc * &self.full.X.fixed_view::<DL, DL>(0, 0) * &h_loc.transpose()
            + h_obs * &noise.Q * &h_obs.transpose();
        self.full
            .X
            .fixed_view_mut::<DF, DF>(fi, fi)
            .copy_from(&feature_cov);
    }

    pub fn forget(&mut self, feature: usize) {
        let fi = DL + feature * DF;
        self.full
            .x
            .rows_mut(fi, DF)
            .iter_mut()
            .for_each(|v| *v = 0.);
        self.full
            .X
            .rows_mut(fi, DF)
            .iter_mut()
            .for_each(|v| *v = 0.);
        self.full
            .X
            .columns_mut(fi, DF)
            .iter_mut()
            .for_each(|v| *v = 0.);
    }
}
