rustsim_spaces/

continuous3d.rs

//! 3D continuous space with spatial-hashing neighbor queries.
//!
//! [`ContinuousSpace3D`] spans `[0, extent_x) x [0, extent_y) x [0, extent_z)`
//! and uses a uniform 3D grid of cells (side length = `spacing`) to accelerate
//! neighbor lookups.
//!
//! Supports periodic (toroidal) and non-periodic (bounded) boundaries.

use rand::Rng;
use rustsim_core::{
    interaction::{PositionedAgent, SpaceInteraction},
    space::Space,
    types::AgentId,
};
use std::collections::HashMap;
use thiserror::Error;

/// 3D position in continuous space.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ContinuousPos3D {
    /// X coordinate.
    pub x: f64,
    /// Y coordinate.
    pub y: f64,
    /// Z coordinate.
    pub z: f64,
}

impl ContinuousPos3D {
    /// Create a new 3D position.
    pub fn new(x: f64, y: f64, z: f64) -> Self {
        Self { x, y, z }
    }

    /// The origin `(0, 0, 0)`.
    pub fn zero() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            z: 0.0,
        }
    }

    /// Euclidean distance to another position (ignoring periodic boundaries).
    pub fn distance_to(&self, other: &Self) -> f64 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        let dz = self.z - other.z;
        (dx * dx + dy * dy + dz * dz).sqrt()
    }

    /// Vector length (magnitude).
    pub fn length(&self) -> f64 {
        (self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
    }

    /// Unit vector in the same direction, or zero if the length is negligible.
    pub fn normalized(&self) -> Self {
        let len = self.length();
        if len < 1e-12 {
            Self::zero()
        } else {
            Self::new(self.x / len, self.y / len, self.z / len)
        }
    }
}

impl std::ops::Add for ContinuousPos3D {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self::new(self.x + rhs.x, self.y + rhs.y, self.z + rhs.z)
    }
}

impl std::ops::Sub for ContinuousPos3D {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self::new(self.x - rhs.x, self.y - rhs.y, self.z - rhs.z)
    }
}

impl std::ops::Mul<f64> for ContinuousPos3D {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self {
        Self::new(self.x * rhs, self.y * rhs, self.z * rhs)
    }
}

/// Errors returned by 3D continuous space configuration validation.
#[derive(Debug, Clone, Copy, PartialEq, Error)]
pub enum ContinuousSpace3DConfigError {
    /// Space extents must be strictly positive.
    #[error("space extents must be positive")]
    InvalidExtent,
    /// Spatial-hash spacing must be strictly positive.
    #[error("spacing must be positive")]
    InvalidSpacing,
}

/// Errors returned by 3D continuous space operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum ContinuousSpace3DError {
    /// The position is outside the space extent (and the space is not periodic).
    #[error("position is out of bounds")]
    OutOfBounds,
}

/// 3D continuous space with uniform-grid spatial hashing for fast neighbor queries.
///
/// The space spans `[0, extent_x) x [0, extent_y) x [0, extent_z)`.
/// An internal grid of cells (side length = `spacing`) accelerates
/// `nearby_ids_euclidean` to O(k) where k is the number of agents in the
/// neighborhood, rather than O(n).
#[derive(Debug, Clone)]
pub struct ContinuousSpace3D {
    extent_x: f64,
    extent_y: f64,
    extent_z: f64,
    periodic: bool,
    spacing: f64,
    grid_w: usize,
    grid_h: usize,
    grid_d: usize,
    cells: Vec<Vec<AgentId>>,
    agent_positions: HashMap<AgentId, ContinuousPos3D>,
}

impl ContinuousSpace3D {
    /// Create a new 3D continuous space.
    ///
    /// # Arguments
    ///
    /// - `extent_x`, `extent_y`, `extent_z` - space dimensions.
    /// - `periodic` - whether boundaries wrap.
    /// - `spacing` - cell side length for the spatial hash grid.
    pub fn new(
        extent_x: f64,
        extent_y: f64,
        extent_z: f64,
        periodic: bool,
        spacing: f64,
    ) -> Result<Self, ContinuousSpace3DConfigError> {
        if extent_x <= 0.0 || extent_y <= 0.0 || extent_z <= 0.0 {
            return Err(ContinuousSpace3DConfigError::InvalidExtent);
        }
        if spacing <= 0.0 {
            return Err(ContinuousSpace3DConfigError::InvalidSpacing);
        }
        let grid_w = (extent_x / spacing).ceil() as usize;
        let grid_h = (extent_y / spacing).ceil() as usize;
        let grid_d = (extent_z / spacing).ceil() as usize;
        let cells = vec![Vec::new(); grid_w * grid_h * grid_d];
        Ok(Self {
            extent_x,
            extent_y,
            extent_z,
            periodic,
            spacing,
            grid_w,
            grid_h,
            grid_d,
            cells,
            agent_positions: HashMap::new(),
        })
    }

    /// Space dimensions as `(extent_x, extent_y, extent_z)`.
    pub fn extent(&self) -> (f64, f64, f64) {
        (self.extent_x, self.extent_y, self.extent_z)
    }

    /// Whether the space wraps around (toroidal boundaries).
    pub fn periodic(&self) -> bool {
        self.periodic
    }

    /// Cell side length for the spatial hash grid.
    pub fn spacing(&self) -> f64 {
        self.spacing
    }

    /// Look up the current position of an agent.
    pub fn agent_position(&self, id: AgentId) -> Option<&ContinuousPos3D> {
        self.agent_positions.get(&id)
    }

    fn normalize_pos(&self, pos: ContinuousPos3D) -> Option<ContinuousPos3D> {
        if self.periodic {
            let x = ((pos.x % self.extent_x) + self.extent_x) % self.extent_x;
            let y = ((pos.y % self.extent_y) + self.extent_y) % self.extent_y;
            let z = ((pos.z % self.extent_z) + self.extent_z) % self.extent_z;
            Some(ContinuousPos3D::new(x, y, z))
        } else if pos.x >= 0.0
            && pos.x < self.extent_x
            && pos.y >= 0.0
            && pos.y < self.extent_y
            && pos.z >= 0.0
            && pos.z < self.extent_z
        {
            Some(pos)
        } else {
            None
        }
    }

    fn pos_to_cell(&self, pos: &ContinuousPos3D) -> (usize, usize, usize) {
        let cx = ((pos.x / self.spacing).floor() as usize).min(self.grid_w - 1);
        let cy = ((pos.y / self.spacing).floor() as usize).min(self.grid_h - 1);
        let cz = ((pos.z / self.spacing).floor() as usize).min(self.grid_d - 1);
        (cx, cy, cz)
    }

    fn cell_index(&self, cx: usize, cy: usize, cz: usize) -> usize {
        (cz * self.grid_h + cy) * self.grid_w + cx
    }

    fn add_to_grid(&mut self, id: AgentId, pos: &ContinuousPos3D) {
        let (cx, cy, cz) = self.pos_to_cell(pos);
        let idx = self.cell_index(cx, cy, cz);
        self.cells[idx].push(id);
    }

    fn remove_from_grid(&mut self, id: AgentId, pos: &ContinuousPos3D) {
        let (cx, cy, cz) = self.pos_to_cell(pos);
        let idx = self.cell_index(cx, cy, cz);
        if let Some(i) = self.cells[idx].iter().position(|v| *v == id) {
            self.cells[idx].swap_remove(i);
        }
    }

    /// Move an agent to a new position, updating the spatial hash.
    ///
    /// Returns the (possibly normalized) final position, or an error if
    /// the position is out of bounds in a non-periodic space.
    pub fn move_agent_pos(
        &mut self,
        id: AgentId,
        new_pos: ContinuousPos3D,
    ) -> Result<ContinuousPos3D, ContinuousSpace3DError> {
        let new_pos = self
            .normalize_pos(new_pos)
            .ok_or(ContinuousSpace3DError::OutOfBounds)?;

        if let Some(old_pos) = self.agent_positions.get(&id).copied() {
            let old_cell = self.pos_to_cell(&old_pos);
            let new_cell = self.pos_to_cell(&new_pos);
            if old_cell != new_cell {
                self.remove_from_grid(id, &old_pos);
                self.add_to_grid(id, &new_pos);
            }
            self.agent_positions.insert(id, new_pos);
        }
        Ok(new_pos)
    }

    fn euclidean_distance_wrapped(&self, a: &ContinuousPos3D, b: &ContinuousPos3D) -> f64 {
        let mut dx = (a.x - b.x).abs();
        let mut dy = (a.y - b.y).abs();
        let mut dz = (a.z - b.z).abs();
        if self.periodic {
            if dx > self.extent_x / 2.0 {
                dx = self.extent_x - dx;
            }
            if dy > self.extent_y / 2.0 {
                dy = self.extent_y - dy;
            }
            if dz > self.extent_z / 2.0 {
                dz = self.extent_z - dz;
            }
        }
        (dx * dx + dy * dy + dz * dz).sqrt()
    }

    /// Find all agent IDs within Euclidean distance `radius` of `pos`.
    ///
    /// Respects periodic boundaries when enabled.
    pub fn nearby_ids_euclidean(&self, pos: &ContinuousPos3D, radius: f64) -> Vec<AgentId> {
        let grid_r = ((radius / self.spacing).ceil() as i32) + 1;
        let (focal_cx, focal_cy, focal_cz) = self.pos_to_cell(pos);
        let mut ids = Vec::new();

        for dz in -grid_r..=grid_r {
            for dy in -grid_r..=grid_r {
                for dx in -grid_r..=grid_r {
                    let nx = focal_cx as i32 + dx;
                    let ny = focal_cy as i32 + dy;
                    let nz = focal_cz as i32 + dz;

                    let cell = if self.periodic {
                        let px =
                            ((nx % self.grid_w as i32) + self.grid_w as i32) % self.grid_w as i32;
                        let py =
                            ((ny % self.grid_h as i32) + self.grid_h as i32) % self.grid_h as i32;
                        let pz =
                            ((nz % self.grid_d as i32) + self.grid_d as i32) % self.grid_d as i32;
                        Some((px as usize, py as usize, pz as usize))
                    } else if nx >= 0
                        && ny >= 0
                        && nz >= 0
                        && (nx as usize) < self.grid_w
                        && (ny as usize) < self.grid_h
                        && (nz as usize) < self.grid_d
                    {
                        Some((nx as usize, ny as usize, nz as usize))
                    } else {
                        None
                    };

                    if let Some((cx, cy, cz)) = cell {
                        let idx = self.cell_index(cx, cy, cz);
                        for &id in &self.cells[idx] {
                            if let Some(agent_pos) = self.agent_positions.get(&id) {
                                if self.euclidean_distance_wrapped(pos, agent_pos) <= radius {
                                    ids.push(id);
                                }
                            }
                        }
                    }
                }
            }
        }
        ids
    }
}

impl Space for ContinuousSpace3D {}

impl<A> SpaceInteraction<A> for ContinuousSpace3D
where
    A: PositionedAgent<Position = ContinuousPos3D>,
{
    type Error = ContinuousSpace3DError;

    fn random_position<R: rand::RngCore>(&self, rng: &mut R) -> A::Position {
        ContinuousPos3D::new(
            rng.gen_range(0.0..self.extent_x),
            rng.gen_range(0.0..self.extent_y),
            rng.gen_range(0.0..self.extent_z),
        )
    }

    fn add_agent(&mut self, agent: &A) -> Result<(), Self::Error> {
        let pos = self
            .normalize_pos(*agent.position())
            .ok_or(ContinuousSpace3DError::OutOfBounds)?;
        let id = agent.id();
        self.agent_positions.insert(id, pos);
        self.add_to_grid(id, &pos);
        Ok(())
    }

    fn remove_agent(&mut self, agent: &A) -> Result<(), Self::Error> {
        let id = agent.id();
        if let Some(pos) = self.agent_positions.remove(&id) {
            self.remove_from_grid(id, &pos);
        }
        Ok(())
    }

    fn nearby_ids(&self, position: &A::Position, radius: usize) -> Vec<AgentId> {
        self.nearby_ids_euclidean(position, radius as f64)
    }
}