use rand::Rng;
use rustsim_core::{
interaction::{PositionedAgent, SpaceInteraction},
space::Space,
types::AgentId,
};
use std::collections::HashMap;
use thiserror::Error;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ContinuousPos {
pub x: f64,
pub y: f64,
}
impl ContinuousPos {
pub fn new(x: f64, y: f64) -> Self {
Self { x, y }
}
pub fn distance_to(&self, other: &Self) -> f64 {
let dx = self.x - other.x;
let dy = self.y - other.y;
(dx * dx + dy * dy).sqrt()
}
pub fn zero() -> Self {
Self { x: 0.0, y: 0.0 }
}
pub fn length(&self) -> f64 {
(self.x * self.x + self.y * self.y).sqrt()
}
pub fn normalized(&self) -> Self {
let len = self.length();
if len < 1e-12 {
Self::zero()
} else {
Self::new(self.x / len, self.y / len)
}
}
}
impl std::ops::Add for ContinuousPos {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self::new(self.x + rhs.x, self.y + rhs.y)
}
}
impl std::ops::Sub for ContinuousPos {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self::new(self.x - rhs.x, self.y - rhs.y)
}
}
impl std::ops::Mul<f64> for ContinuousPos {
type Output = Self;
fn mul(self, rhs: f64) -> Self {
Self::new(self.x * rhs, self.y * rhs)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Error)]
pub enum ContinuousSpaceConfigError {
#[error("space extents must be positive")]
InvalidExtent,
#[error("spacing must be positive")]
InvalidSpacing,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum ContinuousSpaceError {
#[error("position is out of bounds")]
OutOfBounds,
}
#[derive(Debug, Clone)]
pub struct ContinuousSpace2D {
extent_x: f64,
extent_y: f64,
periodic: bool,
spacing: f64,
grid_w: usize,
grid_h: usize,
cells: Vec<Vec<AgentId>>,
agent_positions: HashMap<AgentId, ContinuousPos>,
}
impl ContinuousSpace2D {
pub fn new(
extent_x: f64,
extent_y: f64,
periodic: bool,
spacing: f64,
) -> Result<Self, ContinuousSpaceConfigError> {
if extent_x <= 0.0 || extent_y <= 0.0 {
return Err(ContinuousSpaceConfigError::InvalidExtent);
}
if spacing <= 0.0 {
return Err(ContinuousSpaceConfigError::InvalidSpacing);
}
let grid_w = (extent_x / spacing).ceil() as usize;
let grid_h = (extent_y / spacing).ceil() as usize;
let cells = vec![Vec::new(); grid_w * grid_h];
Ok(Self {
extent_x,
extent_y,
periodic,
spacing,
grid_w,
grid_h,
cells,
agent_positions: HashMap::new(),
})
}
pub fn extent(&self) -> (f64, f64) {
(self.extent_x, self.extent_y)
}
pub fn periodic(&self) -> bool {
self.periodic
}
pub fn spacing(&self) -> f64 {
self.spacing
}
pub fn agent_position(&self, id: AgentId) -> Option<&ContinuousPos> {
self.agent_positions.get(&id)
}
fn normalize_pos(&self, pos: ContinuousPos) -> Option<ContinuousPos> {
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;
Some(ContinuousPos::new(x, y))
} else if pos.x >= 0.0 && pos.x < self.extent_x && pos.y >= 0.0 && pos.y < self.extent_y {
Some(pos)
} else {
None
}
}
fn pos_to_cell(&self, pos: &ContinuousPos) -> (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);
(cx, cy)
}
fn cell_index(&self, cx: usize, cy: usize) -> usize {
cy * self.grid_w + cx
}
fn add_to_grid(&mut self, id: AgentId, pos: &ContinuousPos) {
let (cx, cy) = self.pos_to_cell(pos);
let idx = self.cell_index(cx, cy);
self.cells[idx].push(id);
}
fn remove_from_grid(&mut self, id: AgentId, pos: &ContinuousPos) {
let (cx, cy) = self.pos_to_cell(pos);
let idx = self.cell_index(cx, cy);
if let Some(i) = self.cells[idx].iter().position(|v| *v == id) {
self.cells[idx].swap_remove(i);
}
}
pub fn move_agent_pos(
&mut self,
id: AgentId,
new_pos: ContinuousPos,
) -> Result<ContinuousPos, ContinuousSpaceError> {
let new_pos = self
.normalize_pos(new_pos)
.ok_or(ContinuousSpaceError::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: &ContinuousPos, b: &ContinuousPos) -> f64 {
let mut dx = (a.x - b.x).abs();
let mut dy = (a.y - b.y).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;
}
}
(dx * dx + dy * dy).sqrt()
}
pub fn nearby_ids_euclidean(&self, pos: &ContinuousPos, radius: f64) -> Vec<AgentId> {
let grid_r = ((radius / self.spacing).ceil() as i32) + 1;
let (focal_cx, focal_cy) = self.pos_to_cell(pos);
let mut ids = Vec::new();
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 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;
Some((px as usize, py as usize))
} else if nx >= 0
&& ny >= 0
&& (nx as usize) < self.grid_w
&& (ny as usize) < self.grid_h
{
Some((nx as usize, ny as usize))
} else {
None
};
if let Some((cx, cy)) = cell {
let idx = self.cell_index(cx, cy);
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 ContinuousSpace2D {}
impl<A> SpaceInteraction<A> for ContinuousSpace2D
where
A: PositionedAgent<Position = ContinuousPos>,
{
type Error = ContinuousSpaceError;
fn random_position<R: rand::RngCore>(&self, rng: &mut R) -> A::Position {
ContinuousPos::new(
rng.gen_range(0.0..self.extent_x),
rng.gen_range(0.0..self.extent_y),
)
}
fn add_agent(&mut self, agent: &A) -> Result<(), Self::Error> {
let pos = self
.normalize_pos(*agent.position())
.ok_or(ContinuousSpaceError::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)
}
}