#[path = "../tests/loader.rs"]
mod loader;
use std::hint::black_box;

use cch_rs::cch::*;
use criterion::{Criterion, criterion_group, criterion_main};
use petgraph::algo::{astar, dijkstra};
use rustc_hash::FxHashMap;

use crate::loader::{DimacsGraph, DimacsLoader};

const GR_PATH: &str = "tests/data/USA-road-t.COL.gr";
const CO_PATH: &str = "tests/data/USA-road-d.COL.co";
const TEST_PROFILE: u8 = 0;

pub fn bench_cch_pipeline(c: &mut Criterion) {
    println!("Loading dataset...");
    let graph = DimacsLoader::load_dimacs(GR_PATH, CO_PATH).expect("Failed to load data");
    let mut profile_weights = FxHashMap::default();
    profile_weights.insert(TEST_PROFILE, graph.weights.clone());
    let profile_weights_ref = &profile_weights;
    println!("Initializing CCH Engine...");
    let engine = CchEngine::new(&graph, &profile_weights);
    let start_node = 0u32;
    let end_node = (graph.num_nodes() - 1) as u32;

    let mut group = c.benchmark_group("CCH_Engine_Operations");

    group.bench_function("Engine_Initialization", |b| {
        b.iter(|| CchEngine::new(black_box(&graph), black_box(&profile_weights)));
    });

    group.bench_function("Query_Shortest_Path", |b| {
        b.iter(|| engine.query_path(black_box(TEST_PROFILE), black_box(start_node), black_box(end_node)));
    });

    let mut modified_weights = graph.weights.clone();
    if let Some(w) = modified_weights.first_mut() {
        *w += 10.0;
    }
    let update_edge = 0usize;
    let base_weight = graph.weights[update_edge];

    group.bench_function("Update_Profile_Weights", |b| {
        ittapi::pause();
        ittapi::resume();
        b.iter(|| engine.update_weights(black_box(TEST_PROFILE), black_box(&modified_weights)));
        ittapi::pause();
    });

    let partial_engine = CchEngine::new(&graph, &profile_weights);
    partial_engine.update_weights_partial(TEST_PROFILE, &[(update_edge, base_weight + 10.0)]);
    partial_engine.update_weights_partial(TEST_PROFILE, &[(update_edge, base_weight)]);
    let mut partial_toggle = false;
    group.bench_function("Update_Profile_Weights_Partial", |b| {
        b.iter(|| {
            partial_toggle = !partial_toggle;
            let new_weight = if partial_toggle { base_weight + 10.0 } else { base_weight + 11.0 };
            partial_engine.update_weights_partial(black_box(TEST_PROFILE), black_box(&[(update_edge, new_weight)]));
        });
    });

    let modified_old_nodes = vec![100, 500, 2000, 5000];
    let new_nodes: Vec<u32> = Vec::new();

    group.bench_function("Incremental_Topology_Update", |b| {
        b.iter(|| engine.update_topology(black_box(&graph), black_box(&modified_old_nodes), black_box(&new_nodes), black_box(profile_weights_ref)));
    });

    group.finish();
}

pub fn bench_algorithm_comparison(c: &mut Criterion) {
    let graph = DimacsLoader::load_dimacs(GR_PATH, CO_PATH).expect("Failed to load data");
    let test_profile: u8 = 0;
    let mut profile_weights = FxHashMap::default();
    profile_weights.insert(test_profile, graph.weights.clone());

    let engine = CchEngine::new(&graph, &profile_weights);
    let pet_graph = DimacsLoader::to_petgraph(&graph);

    let start_node = 0u32;
    let end_node = (graph.num_nodes() - 1) as u32;

    let start_idx = pet_graph.node_indices().nth(start_node as usize).unwrap();
    let end_idx = pet_graph.node_indices().nth(end_node as usize).unwrap();

    let mut group = c.benchmark_group("Shortest_Path_Comparison");

    // ---- 1. CCH Query ----
    group.bench_function("CCH_Query", |b| {
        b.iter(|| {
            let res = engine.query_path(black_box(test_profile), black_box(start_node), black_box(end_node));
            black_box(res);
        });
    });

    // ---- 2. Petgraph Dijkstra ----
    group.bench_function("Dijkstra_Petgraph", |b| {
        b.iter(|| {
            let path = dijkstra(&pet_graph, black_box(start_idx), Some(black_box(end_idx)), |e| *e.weight());
            black_box(path);
        });
    });

    // ---- 3. Petgraph A* ----
    group.bench_function("AStar_Petgraph", |b| {
        b.iter(|| {
            let path = astar(
                &pet_graph,
                black_box(start_idx),
                |finish| finish == end_idx,
                |e| *e.weight(),
                |n| euclidean_distance(n.index(), end_node as usize, &graph),
            );
            black_box(path);
        });
    });

    group.finish();
}

criterion_group!(
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = bench_cch_pipeline, bench_algorithm_comparison
);
criterion_main!(benches);

fn euclidean_distance(u: usize, v: usize, graph: &DimacsGraph) -> f32 {
    let dx = graph.xs[u] - graph.xs[v];
    let dy = graph.ys[u] - graph.ys[v];
    (dx * dx + dy * dy).sqrt()
}
