rustsim-core 0.0.1

//! Space interaction API - add, remove, move, and query agents in space.
//!
//! This module provides the ergonomic helper functions that mirror Julia
//! Agents.jl's `model_accessing_API.jl` and `space_interaction_API.jl`:
//!
//! - [`add_agent`], [`add_agent_random`] - insert an agent into both the store and the space.
//! - [`remove_agent`] - remove an agent from both the store and the space.
//! - [`move_agent`] - change an agent's position in the space.
//! - [`nearby_ids`], [`nearby_ids_except`] - query agents near a position.
//! - [`nearby_agents`], [`nearby_agents_except`] - query agent references near a position.
//! - [`random_agent`], [`random_id`] - pick a random agent.
//! - [`all_ids`] - list all agent IDs.
//!
//! These functions operate on [`StandardModel`] and require the agent type to
//! implement [`PositionedAgent`] and the space to implement [`SpaceInteraction`].
//!
//! [`StandardModel`]: crate::standard::StandardModel

use crate::{
    agent::Agent, scheduler::Scheduler, space::Space, standard::StandardModel, store::AgentStore,
    types::AgentId,
};
use rand::seq::SliceRandom;
use thiserror::Error;

/// Errors that can occur during space interaction operations.
#[derive(Debug, Error)]
pub enum InteractionError<E: std::fmt::Debug + std::fmt::Display> {
    /// An agent with this ID already exists in the store.
    #[error("duplicate agent id {0}")]
    DuplicateId(AgentId),
    /// No agent with this ID was found.
    #[error("agent not found: {0}")]
    AgentNotFound(AgentId),
    /// The agent store and spatial index disagree about this agent.
    #[error("agent {0} is missing from the spatial index at its stored position")]
    SpaceIndexMissing(AgentId),
    /// The agent appears more than once in the spatial index at its stored position.
    #[error("agent {0} appears more than once in the spatial index at its stored position")]
    SpaceIndexDuplicate(AgentId),
    /// The underlying space reported an error (e.g. out of bounds).
    #[error("space error: {0}")]
    Space(E),
    /// A rollback failed after a space mutation error, leaving the model in an unknown state.
    #[error("space rollback failed during {operation}: original error: {source}; rollback error: {rollback}")]
    RollbackFailed {
        /// Operation that was being rolled back.
        operation: &'static str,
        /// Original operation error.
        source: E,
        /// Rollback error.
        rollback: E,
    },
}

/// Extension trait for agents that have a spatial position.
///
/// Implement this for agent types used with spatial models (grids, continuous
/// spaces, graphs, etc.). The position type is determined by the space.
pub trait PositionedAgent: Agent {
    /// The position type (e.g. `(usize, usize)` for grids, `ContinuousPos` for continuous space).
    type Position: Clone;

    /// Current position of the agent.
    fn position(&self) -> &Self::Position;

    /// Update the agent's position.
    fn set_position(&mut self, position: Self::Position);
}

/// Trait that spaces implement to support agent lifecycle and neighbor queries.
///
/// Each space type defines its own `Error` type and provides methods to
/// add/remove agents, generate random positions, and find nearby agent IDs.
pub trait SpaceInteraction<A: PositionedAgent>: Space {
    /// Error type for space operations.
    type Error: std::fmt::Debug + std::fmt::Display;

    /// Generate a random valid position within this space.
    fn random_position<R: rand::RngCore>(&self, rng: &mut R) -> A::Position;

    /// Register an agent with the space at its current position.
    fn add_agent(&mut self, agent: &A) -> Result<(), Self::Error>;

    /// Deregister an agent from the space.
    fn remove_agent(&mut self, agent: &A) -> Result<(), Self::Error>;

    /// Return all agent IDs within `radius` of `position`.
    ///
    /// The meaning of `radius` depends on the space: grid cells (Chebyshev),
    /// Euclidean distance (continuous), or graph hops (graph).
    fn nearby_ids(&self, position: &A::Position, radius: usize) -> Vec<AgentId>;
}

/// Add a positioned agent to both the store and the space.
///
/// Returns [`InteractionError::DuplicateId`] if an agent with the same ID
/// already exists, or [`InteractionError::Space`] if the space rejects the position.
pub fn add_agent<S, A, Store, Props, R, Sch>(
    model: &mut StandardModel<S, A, Store, Props, R, Sch>,
    agent: A,
) -> Result<(), InteractionError<S::Error>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    let id = agent.id();
    if model.agents.contains(id) {
        return Err(InteractionError::DuplicateId(id));
    }

    // Register with the space before committing the store insert. If the
    // space rejects the position, the model remains unchanged.
    model
        .space
        .add_agent(&agent)
        .map_err(InteractionError::Space)?;

    model.agents.insert(agent);

    // Update max_id
    if id > model.max_id {
        model.max_id = id;
    }

    Ok(())
}

/// Validate that every positioned agent in the store is present exactly once
/// in the spatial index at its stored position.
pub fn validate_space_index<S, A, Store, Props, R, Sch>(
    model: &StandardModel<S, A, Store, Props, R, Sch>,
) -> Result<(), InteractionError<S::Error>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    for id in model.agents.iter_ids() {
        let Some(agent) = model.agents.get(id) else {
            continue;
        };
        let matches = model
            .space
            .nearby_ids(agent.position(), 0)
            .into_iter()
            .filter(|candidate| *candidate == id)
            .count();
        match matches {
            0 => return Err(InteractionError::SpaceIndexMissing(id)),
            1 => {}
            _ => return Err(InteractionError::SpaceIndexDuplicate(id)),
        }
    }
    Ok(())
}

/// Remove a positioned agent from both the store and the space.
///
/// Returns `Ok(None)` if no agent with this ID was found.
pub fn remove_agent<S, A, Store, Props, R, Sch>(
    model: &mut StandardModel<S, A, Store, Props, R, Sch>,
    id: AgentId,
) -> Result<Option<A>, InteractionError<S::Error>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    // Need to access agent to remove from space
    if let Some(agent_ref) = model.agents.get(id) {
        model
            .space
            .remove_agent(&*agent_ref)
            .map_err(InteractionError::Space)?;
    } else {
        return Ok(None);
    }

    Ok(model.agents.remove(id))
}

/// Move an agent to a new position in the space.
///
/// Returns [`InteractionError::AgentNotFound`] if no agent with this ID was found,
/// or [`InteractionError::Space`] if the space rejects the new position.
pub fn move_agent<S, A, Store, Props, R, Sch>(
    model: &mut StandardModel<S, A, Store, Props, R, Sch>,
    id: AgentId,
    new_position: A::Position,
) -> Result<(), InteractionError<S::Error>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    let mut agent_ref = model
        .agents
        .get_mut(id)
        .ok_or(InteractionError::AgentNotFound(id))?;

    let old_position = agent_ref.position().clone();

    model
        .space
        .remove_agent(&*agent_ref)
        .map_err(InteractionError::Space)?;

    agent_ref.set_position(new_position);

    if let Err(source) = model.space.add_agent(&*agent_ref) {
        agent_ref.set_position(old_position);
        if let Err(rollback) = model.space.add_agent(&*agent_ref) {
            return Err(InteractionError::RollbackFailed {
                operation: "move_agent",
                source,
                rollback,
            });
        }
        return Err(InteractionError::Space(source));
    }

    Ok(())
}

/// Pick a random agent ID from the model.
///
/// Returns `None` if the agent store is empty.
pub fn random_id<S, A, Store, Props, R, Sch>(
    model: &mut StandardModel<S, A, Store, Props, R, Sch>,
) -> Option<AgentId>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    let ids: Vec<AgentId> = model.agents.iter_ids();
    if ids.is_empty() {
        return None;
    }

    let mut rng = model.rng_mut();
    ids.choose(&mut *rng).copied()
}

/// Get a vector of all agent IDs in the model.
pub fn all_ids<S, A, Store, Props, R, Sch>(
    model: &StandardModel<S, A, Store, Props, R, Sch>,
) -> Vec<AgentId>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    model.agents.iter_ids()
}

/// Query agent IDs near a position, using the space's distance metric.
///
/// The meaning of `radius` depends on the space: grid cells (Chebyshev),
/// Euclidean distance (continuous), or graph hops (graph).
pub fn nearby_ids<S, A, Store, Props, R, Sch>(
    model: &StandardModel<S, A, Store, Props, R, Sch>,
    position: &A::Position,
    radius: usize,
) -> Vec<AgentId>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    model.space.nearby_ids(position, radius)
}

/// Query agent references near a position, using the space's distance metric.
///
/// The meaning of `radius` depends on the space: grid cells (Chebyshev),
/// Euclidean distance (continuous), or graph hops (graph).
pub fn nearby_agents<'a, S, A, Store, Props, R, Sch>(
    model: &'a StandardModel<S, A, Store, Props, R, Sch>,
    position: &A::Position,
    radius: usize,
) -> Vec<std::cell::Ref<'a, A>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    model
        .space
        .nearby_ids(position, radius)
        .into_iter()
        .filter_map(|id| model.agents.get(id))
        .collect()
}

/// Query agent IDs near a position, excluding a specific agent.
///
/// The meaning of `radius` depends on the space: grid cells (Chebyshev),
/// Euclidean distance (continuous), or graph hops (graph).
pub fn nearby_ids_except<S, A, Store, Props, R, Sch>(
    model: &StandardModel<S, A, Store, Props, R, Sch>,
    position: &A::Position,
    radius: usize,
    exclude_id: AgentId,
) -> Vec<AgentId>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    model
        .space
        .nearby_ids(position, radius)
        .into_iter()
        .filter(|&id| id != exclude_id)
        .collect()
}

/// Query agent references near a position, excluding a specific agent.
///
/// The meaning of `radius` depends on the space: grid cells (Chebyshev),
/// Euclidean distance (continuous), or graph hops (graph).
pub fn nearby_agents_except<'a, S, A, Store, Props, R, Sch>(
    model: &'a StandardModel<S, A, Store, Props, R, Sch>,
    position: &A::Position,
    radius: usize,
    exclude_id: AgentId,
) -> Vec<std::cell::Ref<'a, A>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    model
        .space
        .nearby_ids(position, radius)
        .into_iter()
        .filter(|&id| id != exclude_id)
        .filter_map(|id| model.agents.get(id))
        .collect()
}

/// Pick a random agent reference from the model.
///
/// Returns `None` if the agent store is empty.
pub fn random_agent<'a, S, A, Store, Props, R, Sch>(
    model: &'a mut StandardModel<S, A, Store, Props, R, Sch>,
) -> Option<std::cell::Ref<'a, A>>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    let ids = model.agents.iter_ids();
    if ids.is_empty() {
        return None;
    }

    let mut rng = model.rng_mut();
    let id = *ids.choose(&mut *rng)?;
    drop(rng);

    model.agents.get(id)
}

/// Add an agent to the model at a random position in the space.
///
/// The agent must not already exist in the store. If the space rejects the
/// position, returns [`InteractionError::Space`].
pub fn add_agent_random<S, A, Store, Props, R, Sch>(
    model: &mut StandardModel<S, A, Store, Props, R, Sch>,
    mut agent: A,
) -> Result<(), InteractionError<S::Error>>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: rand::RngCore,
    Sch: Scheduler<StandardModel<S, A, Store, Props, R, Sch>>,
{
    let mut rng = model.rng_mut();
    let position = model.space.random_position(&mut *rng);
    drop(rng);

    agent.set_position(position);
    add_agent(model, agent)
}