#[path = "shared.rs"]
mod shared;
use std::hint::black_box;

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

const BENCH_RANGE_RADIUS: f32 = 30.0;

// Configure Criterion with our benchmark timeout.
pub fn configure_criterion() -> Criterion {
    Criterion::default().measurement_time(BENCH_TIMEOUT)
}

/// A generic helper function for range search benchmarks.
///
/// The lifetime `'a` ties the lifetime of the tree reference and the return
/// value. The closure `search_fn` must return a value whose lifetime is at
/// least `'a`.
fn bench_range_search<'a, T, Q, R>(
    name: &str,
    tree: &'a T,
    query: &Q,
    search_fn: impl Fn(&'a T, &Q, f32) -> R,
    cc: &mut Criterion,
) where
    R: 'a,
{
    cc.bench_function(name, |b| {
        b.iter(|| {
            let res = search_fn(tree, query, BENCH_RANGE_RADIUS);
            black_box(res)
        })
    });
}

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();
    bench_range_search(
        "range_rstartree_bbox_2d",
        &tree,
        &query_rect,
        |t, q, _| t.range_search_bbox(q),
        &mut cc,
    );
}

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