rustsim-core 0.0.1

//! Discrete-time standard model.
//!
//! [`StandardModel`] is the primary model type, mirroring Julia Agents.jl
//! `StandardABM`. Each tick:
//!
//! 1. The scheduler produces an ordered list of agent IDs.
//! 2. For each agent, the `agent_step` function is called with a [`StepContext`].
//! 3. Deferred actions (add/remove) are applied after each agent step.
//! 4. The `model_step` function is called once (if provided).
//! 5. Time advances by 1.
//!
//! The `agents_first` flag controls whether agent steps run before or after
//! the model step (default: agents first).
//!
//! # Determinism
//!
//! `StandardModel` is replayable for a fixed seed when all of the following are
//! held constant:
//! - the same crate version and user step code
//! - the same initial agents and properties
//! - the same scheduler semantics
//! - the same store-dependent ID enumeration behavior where relevant
//! - the same sequence of deferred insert/remove actions
//!
//! In practice:
//! - [`crate::scheduler::ById`] provides the strongest built-in ordering guarantee
//! - [`crate::scheduler::Fastest`] inherits store iteration order
//! - random schedulers remain reproducible only if the pre-randomization ID list is deterministic
//!
//! [`StepContext`]: crate::step_context::StepContext

use crate::{
    agent::Agent,
    interaction::{InteractionError, PositionedAgent, SpaceInteraction},
    model::Model,
    scheduler::Scheduler,
    space::Space,
    store::AgentStore,
    types::{AgentId, Time},
};
use rand::RngCore;
use tracing::{debug, trace};

use std::cell::{Ref, RefCell, RefMut};
use std::marker::PhantomData;

/// Trait for types that can provide a snapshot of current agent IDs.
///
/// Implemented by [`StandardModel`] to allow schedulers to read agent IDs
/// without full `Model` bounds.
pub trait HasAgentIds {
    /// Collect all agent IDs into a new `Vec`.
    fn agent_ids(&self) -> Vec<AgentId>;

    /// Append all agent IDs into `buf` (does not clear it first).
    fn agent_ids_into(&self, buf: &mut Vec<AgentId>);
}

impl<S, A, Store, Props, R, Sch> StandardModel<S, A, Store, Props, R, Sch>
where
    A: PositionedAgent,
    S: SpaceInteraction<A>,
    Store: AgentStore<A>,
    R: RngCore,
    Sch: Scheduler<Self>,
{
    /// Insert a positioned agent into both the store and the space atomically.
    pub fn insert_positioned_agent(&mut self, agent: A) -> Result<(), InteractionError<S::Error>> {
        crate::interaction::add_agent(self, agent)
    }

    /// Remove a positioned agent from both the store and the space atomically.
    pub fn remove_positioned_agent(
        &mut self,
        id: AgentId,
    ) -> Result<Option<A>, InteractionError<S::Error>> {
        crate::interaction::remove_agent(self, id)
    }

    /// Move a positioned agent, updating both the agent value and spatial index.
    pub fn move_positioned_agent(
        &mut self,
        id: AgentId,
        new_position: A::Position,
    ) -> Result<(), InteractionError<S::Error>> {
        crate::interaction::move_agent(self, id, new_position)
    }

    /// Validate that all stored agents are represented by the spatial index.
    pub fn validate_space_index(&self) -> Result<(), InteractionError<S::Error>> {
        crate::interaction::validate_space_index(self)
    }

    /// Advance one tick while applying deferred add/remove actions to both the
    /// agent store and the spatial index.
    pub fn step_spatial(&mut self) -> Result<(), InteractionError<S::Error>> {
        let has_agent_step = self.agent_step_ctx.is_some();
        let has_model_step = self.model_step.is_some();

        if !(has_agent_step || has_model_step) {
            trace!("spatial step skipped: no agent_step or model_step defined");
            return Ok(());
        }

        if self.agents_first {
            self.step_agents_spatial()?;
            self.step_model();
        } else {
            self.step_model();
            self.step_agents_spatial()?;
        }

        self.time = match self.time {
            Time::Discrete(t) => Time::Discrete(t.saturating_add(1)),
            Time::Continuous(t) => Time::Continuous(t + 1.0),
        };

        debug!(time = %self.time, agents = self.agents.len(), "spatial step completed");
        Ok(())
    }

    /// Advance `n` spatially-consistent ticks.
    pub fn run_spatial(&mut self, steps: usize) -> Result<(), InteractionError<S::Error>> {
        for _ in 0..steps {
            self.step_spatial()?;
        }
        Ok(())
    }

    fn apply_deferred_actions_spatial(
        &mut self,
        deferred: &mut Vec<crate::step_context::DeferredAction<A>>,
    ) -> Result<(), InteractionError<S::Error>> {
        use crate::step_context::DeferredAction;

        if deferred.is_empty() {
            return Ok(());
        }

        for action in deferred.drain(..) {
            match action {
                DeferredAction::RemoveAgent(id) => {
                    crate::interaction::remove_agent(self, id)?;
                }
                DeferredAction::InsertAgent(agent) => {
                    crate::interaction::add_agent(self, agent)?;
                }
            }
        }
        Ok(())
    }

    fn step_agents_spatial(&mut self) -> Result<(), InteractionError<S::Error>> {
        let mut ids = std::mem::take(&mut self.schedule_buf);
        ids.clear();
        {
            let mut sched = self.scheduler.borrow_mut();
            sched.schedule_into(&*self, &mut ids);
        }

        if self.agent_step_ctx.is_none() {
            self.schedule_buf = ids;
            return Ok(());
        }

        let mut deferred = std::mem::take(&mut self.deferred_buf);
        deferred.clear();

        for &id in &ids {
            let Some(mut agent_ref) = self.agents.get_mut(id) else {
                continue;
            };

            {
                let mut rng = self.rng.borrow_mut();
                let mut sched = self.scheduler.borrow_mut();

                let mut ctx = crate::step_context::StepContext {
                    space: &mut self.space,
                    properties: &mut self.properties,
                    rng: &mut *rng,
                    scheduler: &mut *sched,
                    deferred: &mut deferred,
                    _agent: PhantomData,
                };

                if let Some(step_fn) = self.agent_step_ctx.as_mut() {
                    step_fn(&mut *agent_ref, &mut ctx);
                }
            }

            drop(agent_ref);
            self.apply_deferred_actions_spatial(&mut deferred)?;
        }

        self.deferred_buf = deferred;
        self.schedule_buf = ids;
        Ok(())
    }
}

/// Type alias for the boxed agent step function.
///
/// The step function receives a mutable reference to the current agent and
/// a [`StepContext`] providing safe access to space, properties, and RNG.
///
/// [`StepContext`]: crate::step_context::StepContext
pub type AgentStepFn<S, A, Props, R, Sch> =
    Box<dyn for<'a> FnMut(&mut A, &mut crate::step_context::StepContext<'a, S, A, Props, R, Sch>)>;

/// Discrete-time agent-based model.
///
/// This is the main model type for tick-based simulations. It stores agents
/// in an [`AgentStore`], uses a [`Scheduler`] to determine activation order,
/// and calls user-provided step functions each tick.
///
/// # Type Parameters
///
/// - `S` - space type (e.g. `NothingSpace`, `Grid2D`)
/// - `A` - agent type implementing [`Agent`]
/// - `Store` - agent container implementing [`AgentStore<A>`]
/// - `Props` - user-defined model properties (use `()` if unused)
/// - `R` - RNG type (e.g. `StdRng`)
/// - `Sch` - scheduler type (e.g. [`Fastest`], [`Randomly`])
///
/// # Example
///
/// ```ignore
/// type MyModel = StandardModel<NothingSpace, Particle, HashMapStore<Particle>, (), StdRng, Fastest>;
///
/// let mut model = MyModel::new(store, space, Fastest::new(), (), rng, Some(Box::new(step_fn)), None, true);
/// model.step_n(100);
/// ```
///
/// [`Agent`]: crate::agent::Agent
/// [`AgentStore`]: crate::store::AgentStore
/// [`Fastest`]: crate::scheduler::Fastest
/// [`Randomly`]: crate::scheduler::Randomly
pub struct StandardModel<S, A, Store, Props, R, Sch>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: RngCore,
{
    pub(crate) agents: Store,
    pub(crate) space: S,
    pub(crate) scheduler: RefCell<Sch>,
    pub(crate) properties: Props,
    pub(crate) rng: RefCell<R>,
    pub(crate) time: Time,
    pub(crate) max_id: AgentId,
    pub(crate) agents_first: bool,
    pub(crate) agent_step_ctx: Option<AgentStepFn<S, A, Props, R, Sch>>,
    pub(crate) model_step: Option<fn(&mut Self)>,
    pub(crate) deferred_buf: Vec<crate::step_context::DeferredAction<A>>,
    pub(crate) schedule_buf: Vec<AgentId>,
    pub(crate) agent_marker: PhantomData<A>,
}

impl<S, A, Store, Props, R, Sch> StandardModel<S, A, Store, Props, R, Sch>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: RngCore,
    Sch: Scheduler<Self>,
{
    #[allow(clippy::too_many_arguments)]
    /// Create a new `StandardModel`.
    ///
    /// # Arguments
    ///
    /// - `agents` - pre-populated agent store.
    /// - `space` - the simulation space.
    /// - `scheduler` - controls agent activation order.
    /// - `properties` - user-defined model properties.
    /// - `rng` - seeded random number generator.
    /// - `agent_step_ctx` - optional per-agent step function.
    /// - `model_step` - optional per-tick model step function.
    /// - `agents_first` - if `true`, agent steps run before the model step.
    pub fn new(
        agents: Store,
        space: S,
        scheduler: Sch,
        properties: Props,
        rng: R,
        agent_step_ctx: Option<AgentStepFn<S, A, Props, R, Sch>>,
        model_step: Option<fn(&mut Self)>,
        agents_first: bool,
    ) -> Self {
        let max_id = agents.iter_ids().into_iter().max().unwrap_or(0);
        Self {
            agents,
            space,
            scheduler: RefCell::new(scheduler),
            properties,
            rng: RefCell::new(rng),
            time: Time::Discrete(0),
            max_id,
            agents_first,
            agent_step_ctx,
            model_step,
            deferred_buf: Vec::new(),
            schedule_buf: Vec::new(),
            agent_marker: PhantomData,
        }
    }

    /// Create a model with no step functions configured yet.
    ///
    /// This is the most ergonomic starting point for external consumers who
    /// want to configure the model via builder-style methods.
    ///
    /// Defaults:
    /// - no `agent_step_ctx`
    /// - no `model_step`
    /// - `agents_first = true`
    pub fn new_base(agents: Store, space: S, scheduler: Sch, properties: Props, rng: R) -> Self {
        Self::new(agents, space, scheduler, properties, rng, None, None, true)
    }

    /// Create a model with an agent step configured and no model step.
    ///
    /// This avoids `Some(Box::new(...))` boilerplate in the common case of a
    /// purely agent-driven simulation.
    pub fn new_with_agent_step(
        agents: Store,
        space: S,
        scheduler: Sch,
        properties: Props,
        rng: R,
        agent_step_ctx: impl for<'a> FnMut(&mut A, &mut crate::step_context::StepContext<'a, S, A, Props, R, Sch>)
            + 'static,
        agents_first: bool,
    ) -> Self {
        Self::new(
            agents,
            space,
            scheduler,
            properties,
            rng,
            Some(Box::new(agent_step_ctx)),
            None,
            agents_first,
        )
    }

    /// Create a model with a model step configured and no agent step.
    pub fn new_with_model_step(
        agents: Store,
        space: S,
        scheduler: Sch,
        properties: Props,
        rng: R,
        model_step: fn(&mut Self),
        agents_first: bool,
    ) -> Self {
        Self::new(
            agents,
            space,
            scheduler,
            properties,
            rng,
            None,
            Some(model_step),
            agents_first,
        )
    }

    /// Builder-style method to set or replace the agent step function.
    pub fn with_agent_step_ctx(
        mut self,
        agent_step_ctx: impl for<'a> FnMut(&mut A, &mut crate::step_context::StepContext<'a, S, A, Props, R, Sch>)
            + 'static,
    ) -> Self {
        self.agent_step_ctx = Some(Box::new(agent_step_ctx));
        self
    }

    /// Builder-style method to set or replace the model step function.
    pub fn with_model_step(mut self, model_step: fn(&mut Self)) -> Self {
        self.model_step = Some(model_step);
        self
    }

    /// Builder-style method to configure whether agent steps run before model steps.
    pub fn with_agents_first(mut self, agents_first: bool) -> Self {
        self.agents_first = agents_first;
        self
    }

    /// Current simulation time.
    pub fn time(&self) -> Time {
        self.time
    }

    /// Mutable access to the model's RNG (via `RefCell`).
    pub fn rng_mut(&self) -> RefMut<'_, R> {
        self.rng.borrow_mut()
    }

    /// Immutable reference to the simulation space.
    pub fn space(&self) -> &S {
        &self.space
    }

    /// Mutable reference to the simulation space.
    pub fn space_mut(&mut self) -> &mut S {
        &mut self.space
    }

    /// Immutable reference to user-defined properties.
    pub fn properties(&self) -> &Props {
        &self.properties
    }

    /// Mutable reference to user-defined properties.
    pub fn properties_mut(&mut self) -> &mut Props {
        &mut self.properties
    }

    /// Borrow an agent immutably by ID.
    pub fn agent(&self, id: AgentId) -> Option<Ref<'_, A>> {
        self.agents.get(id)
    }

    /// Borrow an agent mutably by ID.
    pub fn agent_mut(&self, id: AgentId) -> Option<RefMut<'_, A>> {
        self.agents.get_mut(id)
    }

    /// Insert an agent into the store.
    ///
    /// Returns `Err(agent)` if an agent with the same ID already exists
    /// (the agent is returned back to the caller).
    pub fn insert_agent(&mut self, agent: A) -> Result<(), A> {
        let id = agent.id();
        if self.agents.contains(id) {
            return Err(agent);
        }
        self.agents.insert(agent);
        if id > self.max_id {
            self.max_id = id;
        }
        Ok(())
    }

    /// Remove an agent by ID, returning it if found.
    pub fn remove_agent(&mut self, id: AgentId) -> Option<A> {
        self.agents.remove(id)
    }

    /// Generate the next unused agent ID (monotonically increasing).
    pub fn next_id(&mut self) -> AgentId {
        self.max_id += 1;
        self.max_id
    }

    /// Advance the simulation by one time step.
    ///
    /// Runs agent steps and model step in the order determined by `agents_first`,
    /// then increments time by 1.
    pub fn step(&mut self) {
        let has_agent_step = self.agent_step_ctx.is_some();
        let has_model_step = self.model_step.is_some();

        if !(has_agent_step || has_model_step) {
            trace!("step skipped: no agent_step or model_step defined");
            return;
        }

        if self.agents_first {
            self.step_agents();
            self.step_model();
        } else {
            self.step_model();
            self.step_agents();
        }

        self.time = match self.time {
            Time::Discrete(t) => Time::Discrete(t.saturating_add(1)),
            Time::Continuous(t) => Time::Continuous(t + 1.0),
        };

        debug!(time = %self.time, agents = self.agents.len(), "step completed");
    }

    /// Advance the simulation by `n` time steps.
    pub fn step_n(&mut self, n: usize) {
        for _ in 0..n {
            self.step();
        }
    }

    /// Alias for [`step_n`](Self::step_n).
    pub fn run(&mut self, steps: usize) {
        self.step_n(steps);
    }

    fn apply_deferred_actions(
        &mut self,
        deferred: &mut Vec<crate::step_context::DeferredAction<A>>,
    ) {
        use crate::step_context::DeferredAction;

        if deferred.is_empty() {
            return;
        }

        for action in deferred.drain(..) {
            match action {
                DeferredAction::RemoveAgent(id) => {
                    self.remove_agent(id);
                }
                DeferredAction::InsertAgent(agent) => {
                    let _ = self.insert_agent(agent);
                }
            }
        }
    }

    fn step_agents(&mut self) {
        let mut ids = std::mem::take(&mut self.schedule_buf);
        ids.clear();
        {
            let mut sched = self.scheduler.borrow_mut();
            sched.schedule_into(&*self, &mut ids);
        }

        if self.agent_step_ctx.is_none() {
            self.schedule_buf = ids;
            return;
        }

        let mut deferred = std::mem::take(&mut self.deferred_buf);
        deferred.clear();

        for &id in &ids {
            let Some(mut agent_ref) = self.agents.get_mut(id) else {
                continue;
            };

            {
                let mut rng = self.rng.borrow_mut();
                let mut sched = self.scheduler.borrow_mut();

                let mut ctx = crate::step_context::StepContext {
                    space: &mut self.space,
                    properties: &mut self.properties,
                    rng: &mut *rng,
                    scheduler: &mut *sched,
                    deferred: &mut deferred,
                    _agent: PhantomData,
                };

                if let Some(step_fn) = self.agent_step_ctx.as_mut() {
                    step_fn(&mut *agent_ref, &mut ctx);
                }
            }

            drop(agent_ref);

            self.apply_deferred_actions(&mut deferred);
        }

        self.deferred_buf = deferred;
        self.schedule_buf = ids;
    }

    fn step_model(&mut self) {
        if let Some(step_fn) = self.model_step {
            step_fn(self);
        }
    }

    /// Iterator over all agents in the store.
    ///
    /// Returns borrowed references. The order depends on the store implementation
    /// and must not be treated as a stable replay contract unless the store
    /// explicitly guarantees it.
    pub fn agents(&self) -> impl Iterator<Item = Ref<'_, A>> {
        self.agents
            .iter_ids()
            .into_iter()
            .filter_map(|id| self.agents.get(id))
    }

    /// Number of agents currently in the store.
    pub fn agents_len(&self) -> usize {
        self.agents.len()
    }
}

impl<S, A, Store, Props, R, Sch> HasAgentIds for StandardModel<S, A, Store, Props, R, Sch>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: RngCore,
{
    fn agent_ids(&self) -> Vec<AgentId> {
        self.agents.iter_ids()
    }

    fn agent_ids_into(&self, buf: &mut Vec<AgentId>) {
        self.agents.iter_ids_into(buf);
    }
}

impl<S, A, Store, Props, R, Sch> Model for StandardModel<S, A, Store, Props, R, Sch>
where
    A: Agent,
    S: Space,
    Store: AgentStore<A>,
    R: rand::RngCore,
{
    type Agent = A;
    type Space = S;
    type Properties = Props;
    type Rng = R;

    // Implementing GATs
    type AgentRef<'a>
        = Ref<'a, A>
    where
        Self: 'a;
    type AgentRefMut<'a>
        = RefMut<'a, A>
    where
        Self: 'a;

    fn time(&self) -> Time {
        self.time
    }

    fn rng_mut(&self) -> impl std::ops::DerefMut<Target = Self::Rng> + '_ {
        self.rng.borrow_mut()
    }

    fn space(&self) -> &Self::Space {
        &self.space
    }

    fn properties(&self) -> &Self::Properties {
        &self.properties
    }

    fn properties_mut(&mut self) -> &mut Self::Properties {
        &mut self.properties
    }

    fn agent(&self, id: AgentId) -> Option<Self::AgentRef<'_>> {
        self.agents.get(id)
    }

    fn agent_mut(&self, id: AgentId) -> Option<Self::AgentRefMut<'_>> {
        self.agents.get_mut(id)
    }
}