use bayes_estimate::estimators::sir;
use bayes_estimate::estimators::sir::SampleState;
use bayes_estimate::models::KalmanState;
use bayes_estimate::noise::{CorrelatedNoise, UncorrelatedNoise};
use nalgebra::{Const, DMatrix, DVector, Dyn, SMatrix, SVector};
use rand::RngCore;

use fast_slam::{FastSLAM, Feature};

mod full_kalman_slam;

#[test]
/// Experiment with a one dimensional problem.
///
/// Useful to look at the implication of highly correlated features.
fn oned_experiment() {
    // State size
    const DL: usize = 1; // Location states
    const DF: usize = 1; // Feature states
    const DZ: usize = 1; // Observation states
                         // Number of features in map
    let nm = 2;

    let nsamples = 10000;
    experiment::<DL, DF, DZ>(nsamples, nm);
}

fn experiment<const DL: usize, const DF: usize, const DZ: usize>(nsamples: usize, nm: usize) {
    // Numeric type
    type N = f64;
    // We need random numbers
    let mut rng = rand::rng();

    // Truth mode, location and map features
    let true_loc = SVector::<N, DL>::repeat(20.);
    let true_0 = SVector::<N, DF>::repeat(50.);
    let true_1 = SVector::<N, DF>::repeat(70.);

    // Location prediction model
    let loc_noise: N = 100.;
    // Model for FastSLAM prediction
    let f_sampled_model = |x: &SVector<N, DL>, rng: &mut dyn RngCore| -> SVector<N, DL> {
        SVector::<N, DL>::repeat(x[0])
            + sir::normal_noise_sampler(SVector::<N, DL>::repeat(loc_noise))(rng)
    };
    // Model for full Kalman prediction
    let fx = SMatrix::<N, DL, DL>::identity();
    let f_q = SVector::<N, DL>::repeat(loc_noise.powi(2));

    // Feature observation model, distance from location to feature
    let hx = SMatrix::<N, DZ, DF>::identity();
    let hx_i = SMatrix::<N, DF, DZ>::identity();
    let h = |loc: &SVector<N, DL>, f: &SVector<N, DF>| &hx * (f - SVector::<N, DF>::repeat(loc[0]));

    let obs_noise_0 = SMatrix::<N, DZ, DZ>::identity() * 5.;
    let obs_noise_1 = SMatrix::<N, DZ, DZ>::identity() * 3.;

    // Model for FastSLAM new feature observation
    let feature_model_0 = |z: &SVector<N, DZ>, loc: &SVector<N, DL>| Feature {
        x: &hx_i * z + SVector::<N, DF>::repeat(loc[0]),
        X: &hx_i * &obs_noise_0 * &hx_i.transpose(),
    };
    let feature_model_1 = |z: &SVector<N, DZ>, loc: &SVector<N, DL>| Feature {
        x: &hx_i * z + SVector::<N, DF>::repeat(loc[0]),
        X: &hx_i * &obs_noise_1 * &hx_i.transpose(),
    };

    // Model for FastSLAM feature observation
    let h_innov_model_0 = |z: &SVector<N, DZ>, loc: &SVector<N, DL>, f: &SVector<N, DF>| {
        (
            z - h(loc, f),
            CorrelatedNoise::<N, Const<DZ>> {
                Q: obs_noise_0.clone(),
            },
        )
    };
    let h_innov_model_1 = |z: &SVector<N, DZ>, loc: &SVector<N, DL>, f: &SVector<N, DF>| {
        (
            z - h(loc, f),
            CorrelatedNoise::<N, Const<DZ>> {
                Q: obs_noise_1.clone(),
            },
        )
    };

    // Model for full Kalman feature observation
    let hx_loc = SMatrix::<N, DZ, DL>::repeat(-1.);
    let hx_map = SMatrix::<N, DZ, DF>::identity();
    let h_loc = SMatrix::<N, DF, DL>::repeat(1.);
    let h_obs = SMatrix::<N, DF, DZ>::identity();

    // Set up the initial state and covariance
    // Location with very little uncertainty
    let init_loc = true_loc.clone();
    let mut init_loc_covariance = SMatrix::<N, DL, DL>::zeros();
    init_loc_covariance[(0, 0)] = 0.0010;

    // Fast_SLAM filter
    let mut roughener = |s: &mut sir::Samples<N, Const<DL>>, rng: &mut dyn RngCore| {
        sir::roughen_minmax::<N, Const<DL>>(s, 1., rng)
    };
    let ks = SampleState::from_kalman(KalmanState {
            x: init_loc,
            X: init_loc_covariance,
        }, nsamples, &mut rng)
        .unwrap();
    let mut fast = FastSLAM::<N, Const<DL>, DF>::new(ks);

    // Kalman_SLAM filter
    let full_size = DL + nm * DF;
    let mut kalm = full_kalman_slam::KalmanSLAM::<DL, DF, DZ> {
        full: KalmanState::<N, Dyn> {
            x: DVector::zeros(full_size),
            X: DMatrix::zeros(full_size, full_size),
        },
    };
    // Filter statistics for display
    let mut stat = kalm.full.clone();

    // Initialise location
    kalm.init_location(&KalmanState::<N, Const<DL>> {
        x: init_loc.clone(),
        X: init_loc_covariance.clone(),
    });


    // Observe initial feature 0
    let z0 = h(&true_loc, &true_0) + SVector::<N, DZ>::repeat(0.5); // Observe a relative position between location and map landmark
    fast.observe_new(0, |loc: &SVector<N, DL>| feature_model_0(&z0, loc));
    kalm.observe_new(
        0,
        &z0,
        &h_loc,
        &h_obs,
        &CorrelatedNoise {
            Q: obs_noise_0.clone(),
        },
    );

    // Observe initial feature 1
    let z1 = h(&true_loc, &true_1) + SVector::<N, DZ>::repeat(-1.0);
    fast.observe_new(1, |loc: &SVector<N, DL>| feature_model_1(&z1, loc));
    kalm.observe_new(
        1,
        &z1,
        &h_loc,
        &h_obs,
        &CorrelatedNoise {
            Q: obs_noise_1.clone(),
        },
    );
    fast.statistics_sparse(&mut stat);
    display("Feature Fast", &stat);
    display("Feature Kalm", &kalm.full);

    // Predict the location state forward
    fast.loc.predict_sampled(f_sampled_model, &mut rng);
    kalm.predict(
        &fx,
        &CorrelatedNoise::from_uncorrelated(&UncorrelatedNoise { q: f_q }),
    );
    fast.statistics_sparse(&mut stat);
    display("Predict Fast", &stat);
    display("Predict Kalm", &kalm.full);

    // Observation feature 0
    let z0 = h(&true_loc, &true_0) + SVector::<N, DZ>::repeat(0.5);
    fast.observe(
        0,
        |loc: &SVector<N, DL>, f: &SVector<N, DF>| h_innov_model_0(&z0, loc, f),
        &hx,
    );
    kalm.observe(
        0,
        &z0,
        &hx_loc,
        &hx_map,
        &CorrelatedNoise {
            Q: obs_noise_0.clone(),
        },
    );
    fast.update_resample(&mut sir::standard_resampler, &mut roughener, &mut rng)
        .unwrap();
    fast.statistics_sparse(&mut stat);
    display("ObserveA Fast", &stat);
    display("ObserveA Kalm", &kalm.full);

    // Observation feature 1
    let z1 = h(&true_loc, &true_1) + SVector::<N, DZ>::repeat(1.0);
    fast.observe(
        1,
        |loc: &SVector<N, DL>, f: &SVector<N, DF>| h_innov_model_1(&z1, loc, f),
        &hx,
    );
    kalm.observe(
        1,
        &z1,
        &hx_loc,
        &hx_map,
        &CorrelatedNoise {
            Q: obs_noise_1.clone(),
        },
    );
    fast.update_resample(&mut sir::standard_resampler, &mut roughener, &mut rng)
        .unwrap();
    fast.statistics_sparse(&mut stat);
    display("ObserveB Fast", &stat);
    display("ObserveB Kalm", &kalm.full);

    // Observation feature 0
    let z0 = h(&true_loc, &true_0) + SVector::<N, DZ>::repeat(-0.5);
    fast.observe(
        0,
        |loc: &SVector<N, DL>, f: &SVector<N, DF>| h_innov_model_0(&z0, loc, f),
        &hx,
    );
    kalm.observe(
        0,
        &z0,
        &hx_loc,
        &hx_map,
        &CorrelatedNoise {
            Q: obs_noise_0.clone(),
        },
    );
    fast.update_resample(&mut sir::systematic_resampler, &mut roughener, &mut rng)
        .unwrap();
    fast.statistics_sparse(&mut stat);
    display("ObserveC Fast", &stat);
    display("ObserveC Kalm", &kalm.full);

    // Forget feature 0
    fast.forget(0);
    kalm.forget(0);
    fast.update_resample(&mut sir::systematic_resampler, &mut roughener, &mut rng)
        .unwrap();
    fast.statistics_sparse(&mut stat);
    display("Forget Fast", &stat);
    display("Forget Kalm", &kalm.full);
}

fn display(what: &str, stat: &KalmanState<f64, Dyn>) {
    println!(
        "{} [{}]({:?}) [{},{}]({:?})",
        what,
        stat.x.shape().0,
        stat.x.data.as_vec(),
        stat.X.shape().0,
        stat.X.shape().1,
        stat.X.data.as_vec()
    );
}
