//! Spatial-index benchmarks: bevy observers-style cartesian grid vs the
//! A5-cell index used by `bevy_a5::hash::CellEntityIndex`.
//!
//! Both methods solve "find every entity within R metres of point P" in two
//! steps: bucket entities by a coarse cell key on insert, then query a
//! handful of buckets near P. The benchmark mirrors each method's
//! data-structure layout exactly so the comparison reflects the methods,
//! not implementation accidents:
//!
//! - **cartesian** — `HashMap<(i32, i32), Vec<Entity>>`, fixed-size square
//!   grid. Same shape as `examples/ecs/observers.rs` in the Bevy repo.
//!   Cell size: 200 m, chosen to match the A5 resolution-9 cell area
//!   roughly (so both indexes pay similar bucket-density costs).
//! - **a5** — `HashMap<u64, Vec<Entity>>`, A5 pentagonal cells at
//!   resolution 9. Same layout as `bevy_a5::hash::CellEntityIndex`. Uses
//!   `bevy_a5::query::grid_disk` to enumerate the buckets touched by a
//!   radius-R query.
//!
//! What's deliberately *not* benchmarked: world-spawn time, observer
//! dispatch, change-detection bookkeeping. Those are scheduling concerns,
//! independent of the index itself. The bench measures the cell-arithmetic
//! and bucket-walk only.
//!
//! Run with:
//!   cargo bench --bench spatial_index
//!
//! Or for a fast smoke run:
//!   cargo bench --bench spatial_index -- --sample-size 10 --warm-up-time 1

use bevy_a5::prelude::*;
use bevy_a5::query::grid_disk;
use bevy_ecs::prelude::Entity;
use bevy_ecs::world::World;
use bevy_math::Vec2;
use bevy_platform::collections::HashMap;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
use rand::prelude::*;

const SEED: u64 = 0xBADC_0FFE_E_BAD_BEEF_u64.wrapping_mul(1);
const N_VALUES: &[usize] = &[1_000, 10_000, 100_000];
const QUERY_COUNT: usize = 100;
/// World half-extent (metres) over which entities + queries are scattered.
const SPREAD_M: f32 = 50_000.0;
/// Query radius (metres). 500 m ≈ 2.5 cells at A5 res 9 / 2.5 cells at the
/// 200 m cartesian grid, so both methods walk a comparable bucket
/// neighbourhood.
const QUERY_RADIUS_M: f32 = 500.0;

// -- Cartesian grid ---------------------------------------------------------

/// Mirrors the `SpatialIndex` resource in Bevy's `examples/ecs/observers.rs`,
/// generalised over query radius. Cell size is fixed; queries walk
/// `ceil(radius / cell_size) + 1` rings around the query tile.
const CARTESIAN_CELL_SIZE: f32 = 200.0;

#[derive(Default)]
struct CartesianIndex {
    map: HashMap<(i32, i32), Vec<Entity>>,
}

impl CartesianIndex {
    fn pos_to_tile(pos: Vec2) -> (i32, i32) {
        (
            (pos.x / CARTESIAN_CELL_SIZE).floor() as i32,
            (pos.y / CARTESIAN_CELL_SIZE).floor() as i32,
        )
    }

    fn insert(&mut self, e: Entity, pos: Vec2) {
        let tile = Self::pos_to_tile(pos);
        self.map.entry(tile).or_default().push(e);
    }

    fn nearby(&self, pos: Vec2, radius_m: f32) -> Vec<Entity> {
        let (tx, ty) = Self::pos_to_tile(pos);
        let k = (radius_m / CARTESIAN_CELL_SIZE).ceil() as i32 + 1;
        let mut out = Vec::new();
        for dx in -k..=k {
            for dy in -k..=k {
                if let Some(es) = self.map.get(&(tx + dx, ty + dy)) {
                    out.extend_from_slice(es);
                }
            }
        }
        out
    }
}

// -- A5 cell grid -----------------------------------------------------------

/// A5 resolution. r=9 ≈ 150 m edge cells on Earth, comparable in area to
/// the 200 m cartesian tile.
const A5_RESOLUTION: i32 = 9;
/// All entity positions are projected lat/lon offsets around this anchor —
/// equirectangular treatment is fine because `SPREAD_M = 50 km` is small
/// enough that the projection doesn't drift meaningfully.
const ANCHOR_LAT: f64 = 48.8566;
const ANCHOR_LON: f64 = 2.3522;

fn xy_to_cell(pos: Vec2) -> Option<GeoCell> {
    let metres_per_deg_lon = 111_320.0 * ANCHOR_LAT.to_radians().cos();
    let metres_per_deg_lat = 110_574.0;
    let lon = ANCHOR_LON + (pos.x as f64) / metres_per_deg_lon;
    let lat = ANCHOR_LAT + (pos.y as f64) / metres_per_deg_lat;
    GeoCell::from_lon_lat(lon, lat, A5_RESOLUTION)
}

#[derive(Default)]
struct A5Index {
    map: HashMap<u64, Vec<Entity>>,
}

impl A5Index {
    fn insert(&mut self, e: Entity, pos: Vec2) {
        let Some(cell) = xy_to_cell(pos) else {
            return;
        };
        self.map.entry(cell.raw()).or_default().push(e);
    }

    fn nearby(&self, pos: Vec2, radius_m: f32) -> Vec<Entity> {
        let Some(centre) = xy_to_cell(pos) else {
            return Vec::new();
        };
        // Pick `k` to cover `radius_m`: cell area → characteristic edge,
        // then ceil(r / edge) + 1 rings of margin. (The flock example fixed
        // a similar undercoverage by switching to spherical_cap; here we
        // stay with grid_disk because it parallels the cartesian `k` more
        // directly, and the small query radius keeps the off-axis gap
        // negligible.)
        let cell_edge_m = (centre.area().sqrt()) as f32;
        let k = ((radius_m / cell_edge_m).ceil() as usize).max(1) + 1;
        let cells = grid_disk(&centre, k).unwrap_or_default();
        let mut out = Vec::new();
        for cell in &cells {
            if let Some(es) = self.map.get(&cell.raw()) {
                out.extend_from_slice(es);
            }
        }
        out
    }
}

// -- Fixtures ---------------------------------------------------------------

fn make_entities(world: &mut World, n: usize) -> Vec<Entity> {
    (0..n).map(|_| world.spawn_empty().id()).collect()
}

fn random_positions(n: usize, half_extent: f32, seed: u64) -> Vec<Vec2> {
    let mut rng = StdRng::seed_from_u64(seed);
    (0..n)
        .map(|_| {
            Vec2::new(
                rng.gen_range(-half_extent..half_extent),
                rng.gen_range(-half_extent..half_extent),
            )
        })
        .collect()
}

// -- Benchmarks -------------------------------------------------------------

fn bench_insert(c: &mut Criterion) {
    let mut group = c.benchmark_group("insert");
    for &n in N_VALUES {
        let mut world = World::new();
        let entities = make_entities(&mut world, n);
        let positions = random_positions(n, SPREAD_M, SEED);
        let pairs: Vec<(Entity, Vec2)> = entities.into_iter().zip(positions).collect();

        group.bench_with_input(BenchmarkId::new("cartesian", n), &pairs, |b, pairs| {
            b.iter(|| {
                let mut idx = CartesianIndex::default();
                for (e, p) in pairs {
                    idx.insert(*e, *p);
                }
                black_box(idx)
            });
        });

        group.bench_with_input(BenchmarkId::new("a5", n), &pairs, |b, pairs| {
            b.iter(|| {
                let mut idx = A5Index::default();
                for (e, p) in pairs {
                    idx.insert(*e, *p);
                }
                black_box(idx)
            });
        });
    }
    group.finish();
}

fn bench_query(c: &mut Criterion) {
    let mut group = c.benchmark_group("query");
    for &n in N_VALUES {
        let mut world = World::new();
        let entities = make_entities(&mut world, n);
        let positions = random_positions(n, SPREAD_M, SEED);
        let queries = random_positions(QUERY_COUNT, SPREAD_M, SEED ^ 0xC0FF_EE);

        let mut cartesian = CartesianIndex::default();
        let mut a5 = A5Index::default();
        for (e, p) in entities.iter().zip(&positions) {
            cartesian.insert(*e, *p);
            a5.insert(*e, *p);
        }

        group.bench_with_input(
            BenchmarkId::new("cartesian", n),
            &queries,
            |b, queries| {
                b.iter(|| {
                    for q in queries {
                        black_box(cartesian.nearby(*q, QUERY_RADIUS_M));
                    }
                });
            },
        );

        group.bench_with_input(BenchmarkId::new("a5", n), &queries, |b, queries| {
            b.iter(|| {
                for q in queries {
                    black_box(a5.nearby(*q, QUERY_RADIUS_M));
                }
            });
        });
    }
    group.finish();
}

criterion_group!(benches, bench_insert, bench_query);
criterion_main!(benches);
