use lqr::{compute_target, find_closest_indices};
use nalgebra as na;
use std::error::Error;

#[test]
fn compute_target_1() {
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.1);
    let expected = na::Matrix5x1::new(0.5, 0.0, 0.0, 0.5, 0.05);
    let target = compute_target(&trajectory, (0, 0.5f64), (1, 0.5f64));
    assert_eq!(expected, target);
}

#[test]
fn compute_target_2() {
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.1);
    let expected = na::Matrix5x1::new(0.25, 0.0, 0.0, 0.25, 0.025);
    let target = compute_target(&trajectory, (0, 0.25), (1, 0.75));
    assert_eq!(target, expected);
}

#[test]
fn compute_target_3() {
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.1);
    let expected = na::Matrix5x1::new(0.75, 0.0, 0.0, 0.75, 0.07500000000000001);
    let target = compute_target(&trajectory, (0, 0.75), (1, 0.25));
    assert_eq!(target, expected);
}

#[test]
fn compute_target_4() {
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.1);
    let expected = na::Matrix5x1::new(0.0, 0.0, 0.0, 0.0, 0.0);
    let target = compute_target(&trajectory, (0, 0.0), (1, f64::MAX));
    assert_eq!(target, expected);
}

#[test]
fn index_search() {
    let current_state = na::Matrix5x1::new(0.5, 0.0, 0.0, 0.0, 0.0);
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);

    match find_closest_indices(&current_state, &trajectory) {
        Ok((lower, upper)) => {
            assert_eq!(lower.0, 0);
            assert_eq!(lower.1, 0.5f64);
            assert_eq!(upper.0, 0);
            assert_eq!(upper.1, 0.5f64);
        }
        Err(_) => {
            assert!(false, "failed to find closest index")
        }
    }
}

#[test]
fn index_search_before_path() {
    let current_state = na::Matrix5x1::new(-1.0, 0.0, 0.0, 0.0, 0.0);
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);

    match find_closest_indices(&current_state, &trajectory) {
        Ok((lower, upper)) => {
            assert_eq!(lower.0, 0);
            assert_eq!(lower.1, 1.0);
            assert_eq!(upper.0, 0);
            assert_eq!(upper.1, 1.0);
        }
        Err(_) => {
            assert!(false, "failed to find closest index")
        }
    }
}

#[test]
fn index_search_after_path() {
    let current_state = na::Matrix5x1::new(2.0, 0.0, 0.0, 0.0, 0.0);
    let trajectory = na::Matrix5x2::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);

    match find_closest_indices(&current_state, &trajectory) {
        Ok((lower, upper)) => {
            assert_eq!(lower.0, 1);
            assert_eq!(lower.1, 1.0);
            assert_eq!(upper.0, 1);
            assert_eq!(upper.1, 1.0);
        }
        Err(_) => {
            assert!(false, "failed to find closest index")
        }
    }
}

#[test]
fn lqr_optimization() -> Result<(), Box<dyn Error>> {
    // Define state
    let _x: f64 = 2.0;
    let _y: f64 = 2.0;
    let theta: f64 = 0.34;
    let v: f64 = 3.0;

    // Define controls
    let delta: f64 = 0.0;
    let _acc: f64 = 0.0;

    // wheelbase
    let l = 2.0;

    // compute matrices for the LQR controller
    let a = na::Matrix4::<f64>::new(
        0.0,
        0.0,
        -v * theta.sin(),
        theta.cos(),
        0.0,
        0.0,
        v * theta.cos(),
        theta.sin(),
        0.0,
        0.0,
        0.0,
        delta.tan() / l,
        0.0,
        0.0,
        0.0,
        0.0,
    );
    let b = na::Matrix4x2::<f64>::new(
        0.0,
        0.0,
        0.0,
        0.0,
        v / (l * delta.cos().powf(2.0)),
        0.0,
        0.0,
        1.0,
    );
    let q = na::Matrix4::identity();
    let r = na::Matrix2::identity();

    let mut controller = lqr::LQRController::new()?;
    controller.compute_gain(&a, &b, &q, &r, 1e-12)?;

    // hand calculated results
    let real_gain = na::Matrix2x4::<f64>::new(0.0, 0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 2.0);

    let mut gain_result = controller.k.unwrap().clone();

    // 0 out the gain
    gain_result.apply(|v| {
        if v.abs() < 1e-6 {
            *v = 0.0f64;
        }
    });

    assert_eq!(real_gain, gain_result);

    Ok(())
}