bevy_a5 0.1.1

//! 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)
    }
}

/// System: keep `GeoCell` in sync with `GlobalTransform` for every entity
/// tagged with [`CellTracker`].
///
/// The expensive `GeoCell::from_world_pos` call is skipped when the entity
/// has moved less than `0.5 * sqrt(cell_area_at_resolution)` from the
/// position at which its current cell was last resolved. This is a
/// conservative bound: half the characteristic cell-edge length is always
/// less than the inscribed-circle radius of an A5 pentagon, so we will
/// never miss a cell crossing — only sometimes pay an unnecessary lookup
/// when the entity is 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(
    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;
        }
        let Some(new_cell) = GeoCell::from_world_pos_f32(world, tracker.resolution) else {
            continue;
        };
        cell.set_if_neq(new_cell);
        tracker.last_resolved_at = Some(world);
    }
}

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

    #[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.
        use bevy_ecs::schedule::Schedule;
        use bevy_ecs::world::World;
        use bevy_transform::components::{GlobalTransform, Transform};

        let mut world = World::new();
        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() {
        use bevy_ecs::schedule::Schedule;
        use bevy_ecs::world::World;
        use bevy_transform::components::{GlobalTransform, Transform};

        let mut world = World::new();
        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 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);
            }
        }
    }
}