kiddo 6.0.2

use az::{Az, Cast};
use codspeed_criterion_compat::{
    criterion_group, criterion_main, measurement::WallTime, AxisScale, BatchSize, BenchmarkGroup,
    BenchmarkId, Criterion, PlotConfiguration, Throughput,
};
use kiddo::batch_benches_parameterized;
use kiddo::distance::float::SquaredEuclidean;
use kiddo::{Eytzinger, KdTree, VecOfArenas};
use kiddo::test_utils::{
    build_populated_tree_and_query_points_float,
    process_queries_fixed_parameterized, process_queries_float_parameterized,
};
use kiddo::traits::{Axis, AxisFixed, Content, Index};
use rand::distr::StandardUniform;
use rand_distr::Distribution;

const BUCKET_SIZE: usize = 32;
const QUERY_POINTS_PER_LOOP: usize = 100;
const RADIUS_SMALL: f64 = 0.01;
const RADIUS_MEDIUM: f64 = 0.05;
const RADIUS_LARGE: f64 = 0.25;
type ImmutableTree<A, T, const K: usize, const B: usize> =
    KdTree<A, T, Eytzinger<K>, VecOfArenas<A, T, K, B>, K, B>;

macro_rules! bench_float {
    ($group:ident, $a:ty, $t:ty, $k:tt, $idx: ty, $size:tt, $radius:tt,  $subtype: expr) => {
        bench_query_float::<$a, $t, $k, $idx>(&mut $group, $size, $radius, $subtype);
    };
}

pub fn within_small(c: &mut Criterion) {
    within(c, RADIUS_SMALL, "small");
}

pub fn within_medium(c: &mut Criterion) {
    within(c, RADIUS_MEDIUM, "medium");
}

pub fn within_large(c: &mut Criterion) {
    within(c, RADIUS_LARGE, "large");
}

fn within(c: &mut Criterion, radius: f64, radius_name: &str) {
    let mut group = c.benchmark_group(format!("Query: within, {radius_name} radius"));
    group.throughput(Throughput::Elements(QUERY_POINTS_PER_LOOP as u64));

    let plot_config = PlotConfiguration::default().summary_scale(AxisScale::Logarithmic);
    group.plot_config(plot_config);

    batch_benches_parameterized!(
        group,
        bench_float,
        radius,
        [(f64, 2), (f64, 3), (f64, 4), (f32, 3)],
        [
            (100, u16, u16),
            (1_000, u16, u16),
            (10_000, u16, u16),
            (100_000, u32, u16),
            (1_000_000, u32, u32)
        ]
    );

    group.finish();
}

fn perform_query_float<
    A: Axis,
    T: Content,
    const K: usize,
    const B: usize,
    IDX: Index<T = IDX> + 'static,
>(
    kdtree: &ImmutableTree<A, T, K, BUCKET_SIZE>,
    point: &[A; K],
    radius: f64,
) where
    usize: Cast<IDX>,
    f64: Cast<A>,
{
    let _res = kdtree
        .query(point)
        .within::<SquaredEuclidean<A>>(radius.az::<A>())
        .execute();
}

fn bench_query_float<
    A: Axis + 'static,
    T: Content,
    const K: usize,
    IDX: Index<T = IDX> + 'static,
>(
    group: &mut BenchmarkGroup<WallTime>,
    initial_size: usize,
    radius: f64,
    subtype: &str,
) where
    usize: Cast<IDX>,
    f64: Cast<A>,
    StandardUniform: Distribution<T>,
    StandardUniform: Distribution<[A; K]>,
{
    group.bench_with_input(
        BenchmarkId::new(subtype, initial_size),
        &initial_size,
        |b, &size| {
            b.iter_batched(
                || {
                    build_populated_tree_and_query_points_float::<A, T, K, BUCKET_SIZE, IDX>(
                        size,
                        QUERY_POINTS_PER_LOOP,
                    )
                },
                process_queries_float_parameterized(
                    perform_query_float::<A, T, K, BUCKET_SIZE, IDX>,
                    radius,
                ),
                BatchSize::SmallInput,
            );
        },
    );
}

criterion_group!(benches, within_small, within_medium, within_large);
criterion_main!(benches);