// This is a full-fledged standalone example of an LQR controller used for full trajectory
// tracking for a car (i.e. ackermann vehicle)

// This example loads a trajectory from a csv file which includes x, y positions, yaw (heading) and
// velocity. This trajectory is then tracked and simulated here using the kinematic bicycle model.
// The program finishes successfully when it reaches the end of the trajectory and saves
// a svg file of the reference trajectory and the achieved one by the controller

use lqr::{compute_target, find_closest_indices, LQRController};
use nalgebra as na;
use num_traits::identities::Zero;
use std::env;
use std::error::Error;
use std::ffi::OsString;
use std::fs::File;

// Computes the A state matrix for an LQR controller based on the kinematic bicycle model
fn compute_a_matrix(
    state: &na::Vector4<f64>,
    controls: &na::Vector2<f64>,
    params: &Vec<f64>,
) -> Result<na::Matrix4<f64>, Box<dyn Error>> {
    Ok(na::Matrix4::<f64>::new(
        0.0,
        0.0,
        -state[3] * state[2].sin(),
        state[2].cos(),
        0.0,
        0.0,
        state[3] * state[2].cos(),
        state[2].sin(),
        0.0,
        0.0,
        0.0,
        controls[0].tan() / params[0],
        0.0,
        0.0,
        0.0,
        0.0,
    ))
}

// Computes the B state matrix for an LQR controller based on the kinematic bicycle model
fn compute_b_matrix(
    state: &na::Vector4<f64>,
    controls: &na::Vector2<f64>,
    params: &Vec<f64>,
) -> Result<na::Matrix4x2<f64>, Box<dyn Error>> {
    Ok(na::Matrix4x2::<f64>::new(
        0.0,
        0.0,
        0.0,
        0.0,
        state[3] / (params[0] * controls[0].cos().powf(2.0)),
        0.0,
        0.0,
        1.0,
    ))
}

// Computes the next state based on the current one and the control inputs
fn next_state(
    state: &na::Vector4<f64>,
    controls: &na::Vector2<f64>,
    params: &Vec<f64>,
    dt: f64,
) -> Result<na::Vector4<f64>, Box<dyn Error>> {
    let state_dot = na::Vector4::<f64>::new(
        state[3] * state[2].cos(),
        state[3] * state[2].sin(),
        state[3] * controls[0].tan() / params[0],
        controls[1],
    );
    return Ok(state + state_dot * dt);
}

// Returns the first positional argument sent to this process. If there are no
// positional arguments, then this returns an error.
fn get_first_arg() -> Result<OsString, Box<dyn Error>> {
    match env::args_os().nth(1) {
        None => Err(From::from(
            "expected 1 argument - the filename of the output file",
        )),
        Some(file_path) => Ok(file_path),
    }
}

fn save_data(data: &Vec<Vec<f64>>) -> Result<(), Box<dyn Error>> {
    // Store as CSV first
    let mut wtr = csv::Writer::from_path("lqr_controller_output.csv")?;
    wtr.write_record(&[
        "index", "curr_x", "curr_y", "curr_yaw", "curr_v", "des_x", "des_y", "des_yaw", "des_v",
        "err_x", "err_y", "err_yaw", "err_v", "delta", "acc", "delta_fb", "acc_fb",
    ])?;

    for row in data {
        wtr.serialize(row.as_slice())?;
    }

    wtr.flush()?;

    Ok(())
}

fn main() -> Result<(), Box<dyn Error>> {
    // Load trajectory
    let trajectory_file = get_first_arg()?;
    let file = File::open(trajectory_file)?;
    let mut rdr = csv::Reader::from_reader(file);

    // First load trajectory in a temporary vector
    let mut temp_trajectory = Vec::<f64>::new();
    for result in rdr.records() {
        let record = result?;
        if &record[0] == "path" {
            temp_trajectory.push(record[1].parse::<f64>()?);
            temp_trajectory.push(record[2].parse::<f64>()?);
            // TODO sexier wrap around yaw
            if record[3].parse::<f64>()? < 0.0 {
                temp_trajectory.push(record[3].parse::<f64>()? + 2.0 * std::f64::consts::PI)
            } else {
                temp_trajectory.push(record[3].parse::<f64>()?);
            }
            temp_trajectory.push(record[4].parse::<f64>()?);
        }
    }

    // Convert trajectory to a na matrix
    let trajectory =
        na::OMatrix::<f64, na::Const<4_usize>, na::Const<1000_usize>>::from_vec(temp_trajectory);

    // initialise state and controls
    // State = [ x_position  y_position  theta  velocity ]
    let mut current_state = na::Vector4::<f64>::new(
        trajectory.column(0)[0],
        trajectory.column(0)[1],
        trajectory.column(0)[2],
        trajectory.column(0)[3],
    ); // TODO find sexier initialization

    // Controls = [ delta  acceleration ]
    let mut current_controls = na::Vector2::zero();

    // Model parameters - [ wheelbase ]
    let params = vec![2.637];
    let dt = 0.01;

    // Set tracking costs
    let q = na::Matrix4::from_diagonal(&na::Vector4::new(1.0, 1.0, 1.0, 0.1));
    let r = na::Matrix2::from_diagonal(&na::Vector2::new(1.0, 0.1));

    // Set up controller
    let mut controller = LQRController::new()?;

    let mut tracking_idx = 0; // index of path to target
    let end_idx = trajectory.shape().1 - 1;

    // Pose storage for plotting purposes
    let mut real_traj = Vec::<(f64, f64)>::new();

    // [index, curr_x, curr_y, curr_yaw, curr_v, des_x, des_y, des_yaw, des_v, err_x, err_y, err_yaw, err_v, delta, acc]
    let mut data_store = Vec::<Vec<f64>>::new();

    // Start simulation loop
    while tracking_idx < end_idx {
        // forward propagate dynamics
        current_state = next_state(&current_state, &current_controls, &params, dt)?;
        real_traj.push((current_state[0], current_state[1]));

        // Find where we are along the trajectory and compute our target state
        let closest_indices = find_closest_indices(&current_state, &trajectory);
        let target_state = match closest_indices {
            Ok((lower, upper)) => {
                tracking_idx = lower.0;
                println!(
                    "Tracking at index {} / {}",
                    tracking_idx,
                    trajectory.shape().1
                );
                compute_target(&trajectory, lower, upper)
            }
            Err(_) => {
                println!("Error in finding target state");
                na::Matrix4x1::zeros()
            }
        };

        println!("Current state {}", current_state);
        println!("Desired state {}", target_state);

        let error = target_state - current_state;

        // update matrices
        let a = compute_a_matrix(&current_state, &current_controls, &params)?;
        let b = compute_b_matrix(&current_state, &current_controls, &params)?;

        controller.compute_gain(&a, &b, &q, &r, 1e-6)?;
        let u_feedback = controller.compute_optimal_controls(&current_state, &target_state)?;
        current_controls += u_feedback;

        println!("Controls {}", current_controls);

        // Push data into buffer for dumping later
        // [index, curr_x, curr_y, curr_yaw, curr_v, des_x, des_y, des_yaw, des_v, err_x, err_y, err_yaw, err_v, delta, acc]
        data_store.push(vec![
            tracking_idx as f64,
            current_state[0],
            current_state[1],
            current_state[2],
            current_state[3],
            target_state[0],
            target_state[1],
            target_state[2],
            target_state[3],
            error[0],
            error[1],
            error[2],
            error[3],
            current_controls[0],
            current_controls[1],
            u_feedback[0],
            u_feedback[1],
        ]);
    }

    println!("Reached end of trajectory. Dumping data");
    save_data(&data_store)
}