use nalgebra::{matrix, vector, Const};

use bayes_estimate::estimators::information_root::InformationRootState;
use bayes_estimate::estimators::ud::UDState;
use bayes_estimate::estimators::unscented_duplex::UnscentedDuplexState;
use bayes_estimate::models::KalmanEstimator;
use bayes_estimate::models::{InformationState, KalmanState};

/// Test the initialisation and reconstruction of linear states.
///
/// Tests 'try_from' be reversing the operation and seeing if we get back the linear state we started with.
#[test]
fn test_init_matches_state() {
    // Start with a UD factored state
    let udstate = UDState::new(matrix![1., 3.; 0., 2.], vector![4., 5.]);
    println!("UD {:} {:}", udstate.x, udstate.UD);

    // Kalman state from the UD state
    let kalman_state = udstate.kalman_state().unwrap();
    println!("Kalman {:} {:}", kalman_state.x, kalman_state.X);
    // Should be able to get back the same UD
    assert_ud_eq(&udstate, &UDState::try_from(kalman_state.clone()).unwrap());

    // Information state from Kalman state
    let information_state = InformationState::try_from(kalman_state.clone()).unwrap();
    println!("Information {:} {:}", information_state.i, information_state.I);
    // Should be able to get back the same Kalman state
    assert_kalman_eq(&kalman_state, &information_state.kalman_state().unwrap());

    // Information root state from Kalman state
    let information_root_state = InformationRootState::try_from(kalman_state.clone()).unwrap();
    println!("InformationRoot from Kalman {:} {:}", information_root_state.r, information_root_state.R);
    // Should be able to get back the same Kalman state
    assert_kalman_eq(&kalman_state, &information_root_state.kalman_state().unwrap());
    // Should be also able to get back the same Information state
    assert_information_eq(&information_state, &information_root_state.information_state().unwrap());

    // Information root state from information state
    let information_root_state2 = InformationRootState::try_from(information_state.clone()).unwrap();
    println!("InformationRoot from Information {:} {:}", information_root_state2.r, information_root_state2.R);
    // Should be able to get back the same Kalman state
    assert_kalman_eq(&kalman_state, &information_root_state2.kalman_state().unwrap());
    // Should be also able to get back the same Information state
    assert_information_eq(&information_state, &information_root_state2.information_state().unwrap());

    // Unscented state from Kalman state, just determines the standard Kappa
    let unscented_state = UnscentedDuplexState::try_from(kalman_state.clone()).unwrap();
    println!("Unscented from Kalman {:} {:}", unscented_state.kalman.x, unscented_state.kappa);
    // Kalman state should be copies
    assert_kalman_eq(&kalman_state, &unscented_state.kalman_state().unwrap());
    assert_eq!(unscented_state.kappa, 1.); // standard Kappa = 3 - 2 (dimensions)
}

fn assert_ud_eq(expect: &UDState<f64, Const<2>>, actual: &UDState<f64, Const<2>>) {
    approx::assert_relative_eq!(actual.x[0], expect.x[0], max_relative = 0.00000001);
    approx::assert_relative_eq!(actual.x[1], expect.x[1], max_relative = 0.01);

    approx::assert_abs_diff_eq!(actual.UD[(0, 0)], expect.UD[(0, 0)], epsilon = 0.000000001);
    approx::assert_abs_diff_eq!(actual.UD[(0, 1)], expect.UD[(0, 1)], epsilon = 0.000000001);
    approx::assert_abs_diff_eq!(actual.UD[(1, 1)], expect.UD[(1, 1)], epsilon = 0.000000001);
}


fn assert_kalman_eq(expect: &KalmanState<f64, Const<2>>, actual: &KalmanState<f64, Const<2>>) {
    approx::assert_relative_eq!(actual.x[0], expect.x[0], max_relative = 0.00000001);
    approx::assert_relative_eq!(actual.x[1], expect.x[1], max_relative = 0.01);

    approx::assert_abs_diff_eq!(actual.X[(0, 0)], expect.X[(0, 0)], epsilon = 0.000001);
    approx::assert_abs_diff_eq!(actual.X[(0, 1)], expect.X[(0, 1)], epsilon = 0.000001);
    approx::assert_abs_diff_eq!(actual.X[(0, 1)], expect.X[(1, 0)], epsilon = 0.0000001);
    approx::assert_abs_diff_eq!(actual.X[(1, 1)], expect.X[(1, 1)], epsilon = 0.000001);
}

fn assert_information_eq(
    expect: &InformationState<f64, Const<2>>,
    actual: &InformationState<f64, Const<2>>,
) {
    approx::assert_relative_eq!(actual.i[0], expect.i[0], max_relative = 0.00000001);
    approx::assert_relative_eq!(actual.i[1], expect.i[1], max_relative = 0.01);

    approx::assert_abs_diff_eq!(actual.I[(0, 0)], expect.I[(0, 0)], epsilon = 0.000001);
    approx::assert_abs_diff_eq!(actual.I[(0, 1)], expect.I[(0, 1)], epsilon = 0.000001);
    approx::assert_abs_diff_eq!(actual.I[(0, 1)], expect.I[(1, 0)], epsilon = 0.0000001);
    approx::assert_abs_diff_eq!(actual.I[(1, 1)], expect.I[(1, 1)], epsilon = 0.000001);
}