bevy_a5 0.1.2

//! The [`GeoCell`] component — an A5 pentagonal cell index for positioning entities on a planet.

use bevy_ecs::component::Component;
use bevy_ecs::prelude::*;
use bevy_reflect::Reflect;
use bevy_transform::components::GlobalTransform;

use bevy_math::{DVec3, Vec3};

use a5::{
    cell_area, cell_to_boundary, cell_to_children, cell_to_lonlat, cell_to_parent,
    get_resolution, lonlat_to_cell, LonLat, WORLD_CELL,
};

use crate::coord;

/// A geospatial cell component that wraps an A5 cell index (`u64`).
///
/// This is the planetary equivalent of `big_space::GridCell`. Each entity with a
/// `GeoCell` is positioned at the center of the corresponding A5 pentagonal cell,
/// with its [`Transform`](bevy_transform::components::Transform) providing a
/// local offset within that cell.
///
/// # Example
///
/// ```rust,ignore
/// commands.spawn((
///     GeoCell::from_lon_lat(2.3522, 48.8566, 9),
///     Transform::default(),
/// ));
/// ```
#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Hash, Reflect)]
#[reflect(Component)]
pub struct GeoCell {
    /// The raw A5 cell index.
    cell: u64,
}

impl GeoCell {
    /// Create a [`GeoCell`] from a raw A5 cell index.
    pub fn new(cell: u64) -> Self {
        Self { cell }
    }

    /// Create a [`GeoCell`] from longitude/latitude (in degrees) at the given resolution.
    ///
    /// Returns `None` if the conversion fails.
    pub fn from_lon_lat(longitude: f64, latitude: f64, resolution: i32) -> Option<Self> {
        let lonlat = LonLat::new(longitude, latitude);
        lonlat_to_cell(lonlat, resolution).ok().map(Self::new)
    }

    /// Create a [`GeoCell`] from a world-space 3D position.
    ///
    /// The point is projected radially onto the planet sphere; only its
    /// direction matters. Returns `None` if `pos` is the origin or the cell
    /// lookup fails.
    pub fn from_world_pos(pos: DVec3, resolution: i32) -> Option<Self> {
        coord::world_pos_to_cell(pos, resolution).map(Self::new)
    }

    /// Create a [`GeoCell`] from a single-precision world-space 3D position.
    ///
    /// Convenience wrapper for `from_world_pos(pos.as_dvec3(), ...)`.
    pub fn from_world_pos_f32(pos: Vec3, resolution: i32) -> Option<Self> {
        Self::from_world_pos(pos.as_dvec3(), resolution)
    }

    /// Get the raw `u64` cell index.
    pub fn raw(&self) -> u64 {
        self.cell
    }

    /// Get the center of this cell as a [`LonLat`].
    ///
    /// Returns `None` if the cell index is invalid.
    pub fn center(&self) -> Option<LonLat> {
        cell_to_lonlat(self.cell).ok()
    }

    /// Get the resolution of this cell.
    pub fn resolution(&self) -> i32 {
        get_resolution(self.cell)
    }

    /// Get the parent cell at the next coarser resolution.
    ///
    /// Returns `None` if the cell is already at resolution 0 or the cell is invalid.
    pub fn parent(&self, resolution: i32) -> Option<Self> {
        cell_to_parent(self.cell, Some(resolution))
            .ok()
            .map(Self::new)
    }

    /// Get the immediate parent (one resolution level coarser).
    pub fn parent_immediate(&self) -> Option<Self> {
        cell_to_parent(self.cell, None).ok().map(Self::new)
    }

    /// Get child cells at a finer resolution.
    ///
    /// Returns `None` if the resolution is invalid.
    pub fn children(&self, resolution: i32) -> Option<Vec<Self>> {
        cell_to_children(self.cell, Some(resolution))
            .ok()
            .map(|cells| cells.into_iter().map(Self::new).collect())
    }

    /// Get immediate children (one resolution level finer).
    pub fn children_immediate(&self) -> Option<Vec<Self>> {
        cell_to_children(self.cell, None)
            .ok()
            .map(|cells| cells.into_iter().map(Self::new).collect())
    }

    /// Get the boundary vertices of this cell as a list of [`LonLat`] pairs.
    ///
    /// Returns `None` if the cell is invalid.
    pub fn boundary(&self) -> Option<Vec<LonLat>> {
        cell_to_boundary(self.cell, None).ok()
    }

    /// Get the area of this cell in square meters (on an authalic Earth).
    pub fn area(&self) -> f64 {
        cell_area(self.resolution())
    }

    /// Check if this is the world cell (resolution 0 root).
    pub fn is_world_cell(&self) -> bool {
        self.cell == WORLD_CELL
    }
}

impl Default for GeoCell {
    fn default() -> Self {
        Self { cell: WORLD_CELL }
    }
}

impl From<u64> for GeoCell {
    fn from(cell: u64) -> Self {
        Self::new(cell)
    }
}

impl From<GeoCell> for u64 {
    fn from(cell: GeoCell) -> Self {
        cell.cell
    }
}

impl core::fmt::Display for GeoCell {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "A5Cell(0x{:x}, res={})", self.cell, self.resolution())
    }
}

/// Helper for entities whose world-space position is set by user code (not
/// by the floating-origin recentre system) and which need a `GeoCell` kept
/// in sync with that position.
///
/// Add this component alongside `GeoCell` and a `GlobalTransform`. The
/// [`update_geo_cells_from_world_pos`] system (registered automatically by
/// `BevyA5Plugin`) will rewrite `GeoCell` to match the entity's world
/// position at the configured `resolution`, but uses a position-delta
/// short-circuit to skip the lookup when the entity hasn't moved more than
/// half a cell-edge since the last successful resolve.
///
/// Why this matters: `GeoCell::from_lon_lat` (the underlying `a5::lonlat_to_cell`
/// call) costs ~5–9 µs per call at typical resolutions. For a flock of 10k
/// entities updating every frame, naively recomputing the cell each frame
/// burns ~50–90 ms; in real games most entities don't cross cell boundaries
/// per frame, so the short-circuit avoids ~99 % of those calls.
///
/// ```rust,ignore
/// commands.spawn((
///     GeoCell::default(),
///     CellTracker::new(9), // resolution 9 ≈ 150 m cells on Earth
///     GlobalTransform::default(),
///     Transform::from_xyz(world_x, 0.0, world_z),
/// ));
/// ```
///
/// Pair with [`CellEntityIndex`](crate::hash::CellEntityIndex) for an O(1)
/// "what entities are near here?" lookup that auto-updates as the tracker
/// rewrites cells.
#[derive(Component, Debug, Clone, Reflect)]
#[reflect(Component)]
pub struct CellTracker {
    /// Resolution at which `GeoCell` should be maintained.
    pub resolution: i32,
    /// World-space position the current cell was computed at. Skipping the
    /// lookup is safe as long as the entity stays within
    /// `0.5 * sqrt(cell_area)` of this point — half the characteristic
    /// cell-edge length.
    last_resolved_at: Option<Vec3>,
}

impl CellTracker {
    /// Construct a tracker at the given A5 resolution.
    pub fn new(resolution: i32) -> Self {
        Self {
            resolution,
            last_resolved_at: None,
        }
    }

    /// Reset the cache. The next [`update_geo_cells_from_world_pos`] pass
    /// will perform a full lookup.
    pub fn invalidate(&mut self) {
        self.last_resolved_at = None;
    }
}

impl Default for CellTracker {
    fn default() -> Self {
        Self::new(9)
    }
}

/// Internal cross-entity lookup cache used by
/// [`update_geo_cells_from_world_pos`]. Hashes world positions to a coarse
/// `(i32, i32)` cartesian bucket and remembers which `GeoCell` was last
/// resolved for each bucket. On a hit, the cached cell is verified by
/// comparing the query position against the cell's centre and a
/// conservative inscribed-radius bound — if the verification passes we
/// return the cached cell directly; if not, we fall through to the full
/// `GeoCell::from_world_pos_f32` lookup and refresh the cache.
///
/// The cache is rebuilt whenever the configured resolution changes
/// (typical games stay at one resolution, so this rarely fires). Entry
/// count is capped to keep memory bounded under pathological scatter
/// patterns; on overflow we wipe and rebuild on demand.
///
/// Sphere-3D world positions (FloatingOrigin pattern) get the verified-hit
/// fast path. Flat-2D-projection users will see verifications fail and
/// fall through to the full lookup every time — correct, but no speed-up.
#[doc(hidden)]
#[derive(Default)]
pub struct BucketLookupCache {
    resolution: Option<i32>,
    /// Squared inscribed-radius bound at the current resolution. Cached
    /// once per resolution to avoid recomputing the sqrt each call.
    inscribed_squared: f32,
    /// Side length of one cartesian bucket in metres. Set to ~half the
    /// cell's inscribed radius so adjacent buckets rarely span a cell
    /// boundary.
    bucket_size_m: f64,
    entries: bevy_platform::collections::HashMap<(i32, i32), CachedCell>,
}

#[derive(Clone, Copy)]
struct CachedCell {
    raw: u64,
    /// World-space position of the cell's centre, used to verify each
    /// cache hit by distance comparison.
    centre: Vec3,
}

/// Maximum bucket-cache entries before a wipe-and-rebuild. ~512 KiB at
/// peak: keeps the cache footprint predictable on pathological scatter.
const BUCKET_CACHE_MAX_ENTRIES: usize = 65_536;

impl BucketLookupCache {
    fn ensure_resolution(&mut self, resolution: i32) {
        if self.resolution == Some(resolution) {
            return;
        }
        self.entries.clear();
        self.resolution = Some(resolution);
        let half_edge = (a5::cell_area(resolution) as f32).sqrt() * 0.5;
        self.inscribed_squared = half_edge * half_edge;
        // Bucket smaller than the inscribed radius — most queries hashing
        // to the same bucket will also be in the same cell.
        self.bucket_size_m = (half_edge * 0.5).max(1.0) as f64;
    }

    fn key_for(&self, pos: Vec3) -> (i32, i32) {
        let bs = self.bucket_size_m;
        (
            (pos.x as f64 / bs).floor() as i32,
            (pos.z as f64 / bs).floor() as i32,
        )
    }

    /// Returns `Some(cell)` if the bucket has a cached cell *and* the
    /// query position falls within its inscribed-radius bound; `None`
    /// otherwise. Caller must fall back to a full lookup on `None`.
    fn try_lookup(&self, key: (i32, i32), pos: Vec3) -> Option<GeoCell> {
        let cached = self.entries.get(&key)?;
        if (pos - cached.centre).length_squared() < self.inscribed_squared {
            Some(GeoCell::new(cached.raw))
        } else {
            None
        }
    }

    fn store(&mut self, key: (i32, i32), cell: GeoCell, centre_world: Vec3) {
        if self.entries.len() >= BUCKET_CACHE_MAX_ENTRIES {
            self.entries.clear();
        }
        self.entries.insert(
            key,
            CachedCell {
                raw: cell.raw(),
                centre: centre_world,
            },
        );
    }
}

/// System: keep `GeoCell` in sync with `GlobalTransform` for every entity
/// tagged with [`CellTracker`].
///
/// The expensive `GeoCell::from_world_pos` call is skipped via two layered
/// caches:
///
/// 1. **Per-entity temporal cache** ([`CellTracker`]'s internal
///    `last_resolved_at`) — if
///    the entity hasn't moved further than `0.5 · sqrt(cell_area)` since
///    its previous resolve, the lookup is a no-op. This is the dominant
///    saver in slow-moving game loops.
/// 2. **Cross-entity bucket cache** (`Local<BucketLookupCache>`) — when
///    one entity *does* need to resolve, the result is keyed by a coarse
///    cartesian bucket and reused for any other entity passing through the
///    same neighbourhood that same frame. Verified per-hit so cell
///    boundaries are never crossed silently.
///
/// Both bounds are conservative: half a cell-edge is always less than the
/// inscribed-circle radius of an A5 pentagon, so we never miss a cell
/// crossing — only sometimes pay an unnecessary lookup near a cell
/// boundary.
///
/// Writes `GeoCell` via [`Mut::set_if_neq`] so `Changed<GeoCell>` (and
/// downstream consumers like `CellEntityIndex`) only fires on real cell
/// transitions.
///
/// Registered automatically in `PostUpdate` by `BevyA5Plugin`, before
/// `propagate_geo_transforms`.
pub fn update_geo_cells_from_world_pos(
    planet: Res<crate::planet::PlanetSettings>,
    mut cache: Local<BucketLookupCache>,
    mut q: Query<(&GlobalTransform, &mut GeoCell, &mut CellTracker)>,
) {
    for (gt, mut cell, mut tracker) in q.iter_mut() {
        let world = gt.translation();
        let half_edge = (a5::cell_area(tracker.resolution) as f32).sqrt() * 0.5;
        let needs_resolve = match tracker.last_resolved_at {
            None => true,
            Some(last) => (world - last).length_squared() > half_edge * half_edge,
        };
        if !needs_resolve {
            continue;
        }

        cache.ensure_resolution(tracker.resolution);
        let key = cache.key_for(world);

        let new_cell = if let Some(c) = cache.try_lookup(key, world) {
            c
        } else {
            let Some(c) = GeoCell::from_world_pos_f32(world, tracker.resolution) else {
                continue;
            };
            // Stash the cell's world-space centre for future verification.
            // Skipped silently if the centre projection fails — we still
            // return the correct cell to the caller, just don't cache it.
            if let Some(centre) = coord::cell_to_dvec3(c.raw(), planet.radius) {
                cache.store(key, c, centre.as_vec3());
            }
            c
        };

        cell.set_if_neq(new_cell);
        tracker.last_resolved_at = Some(world);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use bevy_ecs::schedule::Schedule;
    use bevy_ecs::world::World;
    use bevy_transform::components::{GlobalTransform, Transform};

    /// Helper — spin up a minimal World wired to run
    /// `update_geo_cells_from_world_pos` for tests.
    fn world_with_planet() -> World {
        let mut world = World::new();
        world.insert_resource(crate::planet::PlanetSettings::earth());
        world
    }

    #[test]
    fn cell_tracker_skips_resolve_within_half_edge() {
        // When the entity hasn't moved more than half a cell-edge from the
        // last resolve, `update_geo_cells_from_world_pos` must skip the
        // expensive lookup and leave both `last_resolved_at` and `GeoCell`
        // untouched.
        let mut world = world_with_planet();
        let resolution = 9;
        let half_edge = (a5::cell_area(resolution) as f32).sqrt() * 0.5;

        let initial_pos = Vec3::new(6_400_000.0, 0.0, 0.0);
        let initial_cell = GeoCell::from_world_pos_f32(initial_pos, resolution).unwrap();

        let mut tracker = CellTracker::new(resolution);
        tracker.last_resolved_at = Some(initial_pos);

        let nudge = Vec3::new(half_edge * 0.4, 0.0, 0.0);
        let entity = world
            .spawn((
                initial_cell,
                tracker,
                GlobalTransform::from(Transform::from_translation(initial_pos + nudge)),
            ))
            .id();

        let mut schedule = Schedule::default();
        schedule.add_systems(update_geo_cells_from_world_pos);
        schedule.run(&mut world);

        let tracker_after = world.get::<CellTracker>(entity).unwrap();
        assert_eq!(
            tracker_after.last_resolved_at,
            Some(initial_pos),
            "tracker should not have re-resolved after a sub-half-edge nudge"
        );
        let cell_after = *world.get::<GeoCell>(entity).unwrap();
        assert_eq!(cell_after, initial_cell);
    }

    #[test]
    fn cell_tracker_resolves_on_first_pass() {
        let mut world = world_with_planet();
        let pos = Vec3::new(6_400_000.0, 0.0, 0.0);
        let entity = world
            .spawn((
                GeoCell::default(),
                CellTracker::new(7),
                GlobalTransform::from(Transform::from_translation(pos)),
            ))
            .id();

        let mut schedule = Schedule::default();
        schedule.add_systems(update_geo_cells_from_world_pos);
        schedule.run(&mut world);

        let tracker_after = world.get::<CellTracker>(entity).unwrap();
        assert!(tracker_after.last_resolved_at.is_some());
        let cell_after = *world.get::<GeoCell>(entity).unwrap();
        assert_eq!(cell_after.resolution(), 7);
        // And it should match a direct lookup.
        let expected = GeoCell::from_world_pos_f32(pos, 7).unwrap();
        assert_eq!(cell_after, expected);
    }

    #[test]
    fn bucket_cache_returns_same_cell_as_uncached_path() {
        // Two entities at slightly-different positions in the same coarse
        // bucket — first miss populates the cache, second one hits the
        // cache and must return the *same* cell as a direct lookup.
        // Locks in the verification step's correctness.
        let mut world = world_with_planet();
        let resolution = 7;

        // Force last_resolved_at = None so each entity will go through
        // the lookup path on the first schedule run.
        let pos_a = Vec3::new(6_400_000.0, 0.0, 0.0);
        // pos_b is 1 m east of pos_a — same A5 cell at res 7 (~5 km edges)
        // and same coarse bucket.
        let pos_b = pos_a + Vec3::new(1.0, 0.0, 0.0);

        let entity_a = world
            .spawn((
                GeoCell::default(),
                CellTracker::new(resolution),
                GlobalTransform::from(Transform::from_translation(pos_a)),
            ))
            .id();
        let entity_b = world
            .spawn((
                GeoCell::default(),
                CellTracker::new(resolution),
                GlobalTransform::from(Transform::from_translation(pos_b)),
            ))
            .id();

        let mut schedule = Schedule::default();
        schedule.add_systems(update_geo_cells_from_world_pos);
        schedule.run(&mut world);

        let cell_a = *world.get::<GeoCell>(entity_a).unwrap();
        let cell_b = *world.get::<GeoCell>(entity_b).unwrap();
        let direct = GeoCell::from_world_pos_f32(pos_b, resolution).unwrap();
        assert_eq!(cell_a, direct);
        assert_eq!(cell_b, direct);
    }

    #[test]
    fn bucket_cache_invalidates_on_resolution_change() {
        // Switching tracker resolution between schedule runs must rebuild
        // the cache so we don't return a cell at the wrong resolution.
        let mut world = world_with_planet();
        let pos = Vec3::new(6_400_000.0, 0.0, 0.0);

        let entity = world
            .spawn((
                GeoCell::default(),
                CellTracker::new(7),
                GlobalTransform::from(Transform::from_translation(pos)),
            ))
            .id();

        let mut schedule = Schedule::default();
        schedule.add_systems(update_geo_cells_from_world_pos);
        schedule.run(&mut world);
        assert_eq!(world.get::<GeoCell>(entity).unwrap().resolution(), 7);

        // Bump resolution and force a re-resolve.
        {
            let mut tracker = world.get_mut::<CellTracker>(entity).unwrap();
            tracker.resolution = 9;
            tracker.invalidate();
        }
        schedule.run(&mut world);
        assert_eq!(world.get::<GeoCell>(entity).unwrap().resolution(), 9);
    }

    #[test]
    fn children_returns_uniform_resolution() {
        // Locks in the same invariant as `query::grid_disk` /
        // `query::spherical_cap`: a function that takes a cell + target
        // resolution must return cells uniformly at that resolution.
        let world = GeoCell::new(WORLD_CELL);
        for res in [3, 5, 7] {
            let kids = world.children(res).expect("children should succeed");
            assert!(!kids.is_empty(), "no children at res {}", res);
            for c in &kids {
                assert_eq!(c.resolution(), res);
            }
        }
    }
}