use std::hint::black_box;

use autobarts::{geom::Point, spindex::RStarTree};
use bevy::math::Vec2;
use criterion::{Criterion, criterion_group};

//
// Benchmark Parameters
//
pub const BENCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
pub const BENCH_NUM_INSERT: i32 = 50_000;
pub const BENCH_NODE_CAPACITY: usize = 5;

//
// Data Generation Functions (Raw Data)
//
pub fn generate_2d_data() -> Vec<Point> {
    let data: Vec<Point> = (0..BENCH_NUM_INSERT)
        .map(|i| {
            let point = bevy::prelude::Vec2::new(i as f32, i as f32);

            Point {
                point,
                entity: bevy::prelude::Entity::PLACEHOLDER,
            }
        })
        .collect();

    data
}

// Configure Criterion with a timeout for benchmarks
pub fn configure_criterion() -> Criterion {
    Criterion::default()
        .measurement_time(BENCH_TIMEOUT)
        .sample_size(10)
}

const BENCH_RANGE_RADIUS: f32 = 30.0;

fn benchmark_range_rstartree_2d(_c: &mut Criterion) {
    let points = generate_2d_data();
    let mut tree = RStarTree::new(BENCH_NODE_CAPACITY);

    tree.insert_bulk(points.clone());

    let mut cc = configure_criterion();
    let mut idx = 0;
    let len = points.len();
    cc.bench_function("range_rstartree_2d", |b| {
        b.iter(|| {
            let res = tree.range_search(&points[idx], BENCH_RANGE_RADIUS);
            idx = (idx + 1) % len;
            black_box(res)
        })
    });
}

fn benchmark_range_bbox_rstartree_2d(_c: &mut Criterion) {
    let points = generate_2d_data();
    let mut tree = RStarTree::new(BENCH_NODE_CAPACITY);

    tree.insert_bulk(points);

    let query_rect = bevy::math::bounding::Aabb2d::new(
        Vec2::new(35.0 - BENCH_RANGE_RADIUS, 45.0 - BENCH_RANGE_RADIUS),
        4.0 * Vec2::splat(BENCH_RANGE_RADIUS),
    );
    let mut cc = configure_criterion();
    cc.bench_function("range_rstartree_2d", |b| {
        b.iter(|| {
            let res = tree.range_search_bbox(&query_rect);
            black_box(res)
        })
    });
}

criterion_group!(
    benches,
    benchmark_range_rstartree_2d,
    //benchmark_range_bbox_rstartree_2d,
);