use criterion::{criterion_group, criterion_main, Criterion};
use criterion::BenchmarkId;
use criterion::Throughput;

use fast_motion_planning as mp;
use fast_motion_planning::{MPComponent};
use fast_motion_planning::objects::{ObjectList, Object, Bounds};
use fast_motion_planning::utils;

fn generate_problem() -> mp::MPProblem {
    // Define the axis component ranges
    // comp id, start position, pmin, pmax, start velocity, start acceleration
    let x_param = MPComponent::new("x".to_string(), 0.0, 4.0, 4.0, 3.0, 1.0);
    let y_param = MPComponent::new("y".to_string(), 1.5, 2.8, 3.2, 3.0, 0.5);
    let z_param = MPComponent::new("z".to_string(), 1.0, 1.7, 2.3, 1.0, 1.0);

    // Define the objects in the environment
    let mut objects = ObjectList::new();
    objects.push(Object::new(
        Bounds::new(3.75, 4.25),
        Bounds::new(2.3, 2.8),
        Bounds::new(1.7, 2.3)));
    objects.push(Object::new(
        Bounds::new(3.75, 4.25),
        Bounds::new(3.2, 3.7),
        Bounds::new(1.7, 2.3)));

    mp::MPProblem::new(x_param, y_param, z_param, objects)
}

fn generate_problem_from_file(file_path: String) -> mp::MPProblem {
    // Define the axis component ranges
    // comp id, start position, pmin, pmax, start velocity, start acceleration
    let x_param = MPComponent::new("x".to_string(), 0.0, 4.0, 4.0, 3.0, 1.0);
    let y_param = MPComponent::new("y".to_string(), 1.5, 2.8, 3.2, 3.0, 0.5);
    let z_param = MPComponent::new("z".to_string(), 1.0, 1.7, 2.3, 1.0, 1.0);

    let objects = utils::read_objects_from_file(file_path).expect("Cannot read file");
    mp::MPProblem::new(x_param, y_param, z_param, objects)
}

fn bench_solve_1(dt: usize, method: mp::MPMethod) {
    let time_division = dt;
    let epsilon = 0.01;
    let problem = generate_problem();
    match method {
        mp::MPMethod::BRUTEFORCE => mp::BruteForceFastMotionPlanner::new(time_division, epsilon).solve(problem),
        mp::MPMethod::HEURISTIC => mp::HeuristicFastMotionPlanner::new(time_division, epsilon).solve(problem)
    };
}

fn bench_solve_file(file_path: String, dt:usize, method: mp::MPMethod) {
    let time_division = dt;
    let epsilon = 0.01;
    let problem = generate_problem_from_file(file_path);
    match method {
        mp::MPMethod::BRUTEFORCE => mp::BruteForceFastMotionPlanner::new(time_division, epsilon).solve(problem),
        mp::MPMethod::HEURISTIC => mp::HeuristicFastMotionPlanner::new(time_division, epsilon).solve(problem)
    };
}

fn bench_to_first(file_path: String, dt:usize, method: mp::MPMethod) {
    let time_division = dt;
    let epsilon = 0.01;
    let problem = generate_problem_from_file(file_path);
    match method {
        mp::MPMethod::BRUTEFORCE => mp::BruteForceFastMotionPlanner::new(time_division, epsilon).solver(problem).into_iter().next().unwrap(),
        mp::MPMethod::HEURISTIC => mp::HeuristicFastMotionPlanner::new(time_division, epsilon).solver(problem).into_iter().next().unwrap()
    };
}

fn bench_to_random_first(file_path: String, dt:usize, method: mp::MPMethod) {
    let time_division = dt;
    let epsilon = 0.01;
    let problem = generate_problem_from_file(file_path);
    match method {
        mp::MPMethod::BRUTEFORCE => mp::BruteForceFastMotionPlanner::new(time_division, epsilon).random_solver(problem).into_iter().next().unwrap(),
        mp::MPMethod::HEURISTIC => mp::HeuristicFastMotionPlanner::new(time_division, epsilon).random_solver(problem).into_iter().next().unwrap()
    };
}

fn bench_to_first_trajectory(file_path: String, dt:usize, method: mp::MPMethod) {
    let time_division = dt;
    let epsilon = 0.01;
    let time_start = 0.0;
    let num_samples = 100;
    let problem = generate_problem_from_file(file_path);
    match method {
        mp::MPMethod::BRUTEFORCE => mp::BruteForceFastMotionPlanner::new(time_division, epsilon).solver(problem).into_iter().next().unwrap().generate_fourpl_trajectory(time_start, num_samples),
        mp::MPMethod::HEURISTIC => mp::HeuristicFastMotionPlanner::new(time_division, epsilon).solver(problem).into_iter().next().unwrap().generate_fourpl_trajectory(time_start, num_samples)
    };
}

fn criterion_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("solve scenario 1");
    group.sample_size(25);
    for size in [2, 4, 6, 8, 10, 15, 20].iter() {
        group.bench_with_input(BenchmarkId::new("BruteForce", size), size, |b, &size| {
            b.iter(|| bench_solve_1(size, mp::MPMethod::BRUTEFORCE));
        });

        group.bench_with_input(BenchmarkId::new("Heuristic", size), size, |b, &size| {
            b.iter(|| bench_solve_1(size, mp::MPMethod::HEURISTIC));
        });
    }
    group.finish();

    let file_path = "tests/examples/obs10.csv";
    let mut group = c.benchmark_group(format!("solve scenario {}", &file_path));
    group.sample_size(25);

    for size in [2, 4, 6, 8, 10, 15].iter() {
        group.bench_with_input(BenchmarkId::new("BruteForce", size), size, |b, &size| {
            b.iter(|| bench_solve_file(file_path.to_string(), size, mp::MPMethod::BRUTEFORCE));
        });
        group.bench_with_input(BenchmarkId::new("Heuristic", size), size, |b, &size| {
            b.iter(|| bench_solve_file(file_path.to_string(), size, mp::MPMethod::HEURISTIC));
        });
    }
    group.finish();

    let mut group = c.benchmark_group(format!("time to first solution {}", &file_path));
    group.sample_size(25);

    for size in [2, 4, 6, 8, 10, 15].iter() {
        group.bench_with_input(BenchmarkId::new("BruteForce", size), size, |b, &size| {
            b.iter(|| bench_to_first(file_path.to_string(), size, mp::MPMethod::BRUTEFORCE));
        });
        group.bench_with_input(BenchmarkId::new("Heuristic", size), size, |b, &size| {
            b.iter(|| bench_to_first(file_path.to_string(), size, mp::MPMethod::HEURISTIC));
        });
    }
    group.finish();

    let mut group = c.benchmark_group(format!("time to first solution (random) {}", &file_path));
    group.sample_size(25);

    for size in [2, 4, 6, 8, 10, 15].iter() {
        group.bench_with_input(BenchmarkId::new("BruteForce", size), size, |b, &size| {
            b.iter(|| bench_to_random_first(file_path.to_string(), size, mp::MPMethod::BRUTEFORCE));
        });
        group.bench_with_input(BenchmarkId::new("Heuristic", size), size, |b, &size| {
            b.iter(|| bench_to_random_first(file_path.to_string(), size, mp::MPMethod::HEURISTIC));
        });
    }
    group.finish();

    let mut group = c.benchmark_group(format!("time to first trajectory {}", &file_path));
    group.sample_size(25);

    for size in [2, 4, 6, 8, 10, 15].iter() {
        group.bench_with_input(BenchmarkId::new("BruteForce", size), size, |b, &size| {
            b.iter(|| bench_to_first_trajectory(file_path.to_string(), size, mp::MPMethod::BRUTEFORCE));
        });
        group.bench_with_input(BenchmarkId::new("Heuristic", size), size, |b, &size| {
            b.iter(|| bench_to_first_trajectory(file_path.to_string(), size, mp::MPMethod::HEURISTIC));
        });
    }
    group.finish();
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);