//! Benchmark for various NearestNeighbors implementations.

use acap::euclid::Euclidean;
use acap::exhaustive::ExhaustiveSearch;
use acap::NearestNeighbors;

use criterion::{black_box, criterion_group, criterion_main, Criterion};

use std::iter::FromIterator;

type Point = Euclidean<[f32; 3]>;

/// Generates a spiral used as the benchmark data set.
fn spiral() -> Vec<Point> {
    let mut points = Vec::new();

    let size = 1000;
    let turns = 10.0;
    for i in 0..size {
        let y = 2.0 * (i as f32) / (size as f32) - 1.0;
        let m = (1.0 - y * y).sqrt();
        let theta = turns * y * std::f32::consts::PI;
        let (sin, cos) = theta.sin_cos();
        let x = m * cos * cos;
        let z = m * sin * cos;
        points.push(Euclidean([x, y, z]));
    }

    points
}

fn bench_from_iter(c: &mut Criterion) {
    let points = black_box(spiral());

    let mut group = c.benchmark_group("from_iter");
    group.bench_function("ExhaustiveSearch", |b| b.iter(|| ExhaustiveSearch::from_iter(points.clone())));
    group.finish();
}

fn bench_nearest_neighbors(c: &mut Criterion) {
    let points = black_box(spiral());
    let target = black_box(Euclidean([0.0, 0.0, 0.0]));

    let exhaustive = ExhaustiveSearch::from_iter(points.clone());

    let mut nearest = c.benchmark_group("NearestNeighbors::nearest");
    nearest.bench_function("ExhaustiveSearch", |b| b.iter(|| exhaustive.nearest(&target)));
    nearest.finish();

    let mut nearest_within = c.benchmark_group("NearestNeighbors::nearest_within");
    nearest_within.bench_function("ExhaustiveSearch", |b| b.iter(|| exhaustive.nearest_within(&target, 0.1)));
    nearest_within.finish();

    let mut k_nearest = c.benchmark_group("NearestNeighbors::k_nearest");
    k_nearest.bench_function("ExhaustiveSearch", |b| b.iter(|| exhaustive.k_nearest(&target, 3)));
    k_nearest.finish();

    let mut k_nearest_within = c.benchmark_group("NearestNeighbors::k_nearest_within");
    k_nearest_within.bench_function("ExhaustiveSearch", |b| b.iter(|| exhaustive.k_nearest_within(&target, 3, 0.1)));
    k_nearest_within.finish();
}

criterion_group!(benches, bench_from_iter, bench_nearest_neighbors);
criterion_main!(benches);