Mesa: Agent-Based Modeling Framework

# Core agent-based modeling classes
from mesa import Agent, Model
# Note: AgentSet is available via Agent.create_agents() or model.agents

# Data collection utilities
from mesa import DataCollector

# Traditional spatial modeling (legacy)
from mesa import space

# Modern discrete spatial modeling
from mesa import discrete_space

# Batch running for parameter sweeps
from mesa import batch_run

# Experimental features
from mesa import experimental

# Logging utilities
import mesa.mesa_logging

from mesa import Agent, Model
from mesa.discrete_space import Grid

class MyAgent(Agent):
    """Example agent with basic behavior."""
    
    def __init__(self, model):
        super().__init__(model)
        self.energy = 100
    
    def step(self):
        """Agent behavior executed each simulation step."""
        self.energy -= 1
        if self.energy <= 0:
            self.remove()

class MyModel(Model):
    """Example model with agents on a grid."""
    
    def __init__(self, n_agents: int, width: int, height: int):
        super().__init__()
        self.space = Grid(width, height)
        
        # Create agents
        for i in range(n_agents):
            agent = MyAgent(self)
            # Place agent randomly
            self.space.place_agent(agent, self.space.select_random_empty_cell())
    
    def step(self):
        """Execute one step of the model."""
        self.agents.shuffle_do("step")

# Run the model
model = MyModel(n_agents=50, width=20, height=20)
for i in range(100):
    model.step()
    if len(model.agents) == 0:
        break

from mesa import Agent, Model
# AgentSet is accessible through agent creation or model.agents property

class Agent:
    """Base class for model agents with lifecycle management."""
    
    def __init__(self, model: Model, *args, **kwargs) -> None: ...
    def step(self) -> None: ...
    def advance(self) -> None: ...
    def remove(self) -> None: ...
    
    @classmethod
    def create_agents(cls, model: Model, n: int, *args, **kwargs) -> AgentSet[Agent]: ...
    
    @property
    def random(self) -> Random: ...
    
    @property  
    def rng(self) -> np.random.Generator: ...

class AgentSet(MutableSet, Sequence):
    """Collection class for managing ordered sets of agents with advanced operations."""
    
    def __init__(self, agents: Iterable[Agent], random: Random | None = None): ...
    
    def select(self, filter_func: Callable[[Agent], bool] | None = None,
               at_most: int | float = float("inf"), 
               inplace: bool = False,
               agent_type: type[Agent] | None = None) -> AgentSet: ...
    
    def shuffle(self, inplace: bool = False) -> AgentSet: ...
    def sort(self, key: Callable[[Agent], Any] | str, 
             ascending: bool = False, inplace: bool = False) -> AgentSet: ...
    
    def do(self, method: str | Callable, *args, **kwargs) -> AgentSet: ...
    def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet: ...
    def map(self, method: str | Callable, *args, **kwargs) -> list[Any]: ...

class Model:
    """Base class for agent-based models with simulation management."""
    
    def __init__(self, *args: Any, seed: float | None = None,
                 rng: RNGLike | SeedLike | None = None, **kwargs: Any) -> None: ...
    
    def step(self) -> None: ...
    def run_model(self) -> None: ...
    def register_agent(self, agent): ...
    def deregister_agent(self, agent): ...
    
    @property
    def agents(self) -> AgentSet: ...
    
    @property
    def agents_by_type(self) -> dict[type[Agent], AgentSet]: ...

from mesa import DataCollector

class DataCollector:
    """Data collection system for Mesa models with multiple reporter types."""
    
    def __init__(self, 
                 model_reporters=None, 
                 agent_reporters=None,
                 agenttype_reporters=None, 
                 tables=None): ...
    
    def collect(self, model): ...
    def add_table_row(self, table_name, row, ignore_missing=False): ...
    def get_model_vars_dataframe(self): ...
    def get_agent_vars_dataframe(self): ...
    def get_agenttype_vars_dataframe(self, agent_type): ...
    def get_table_dataframe(self, table_name): ...

# Legacy spatial system (maintenance mode)
from mesa.space import MultiGrid, SingleGrid, ContinuousSpace, NetworkGrid

# Modern discrete spatial system (active development)
from mesa.discrete_space import Grid, HexGrid, Network, Cell, CellAgent

# Property layers for spatial data
from mesa.discrete_space import PropertyLayer
from mesa.space import PropertyLayer as LegacyPropertyLayer

class Grid:
    """Modern grid-based discrete space with cell-centric approach."""
    
    def __init__(self, width: int, height: int, torus: bool = False): ...
    def place_agent(self, agent, cell): ...
    def move_agent(self, agent, cell): ...
    def select_random_empty_cell(self): ...

class CellAgent(Agent):
    """Agent class that understands cell-based spatial interactions."""
    
    def __init__(self, model, cell=None): ...
    def move_to(self, cell): ...
    def get_neighbors(self, include_center=False): ...

from mesa import batch_run

def batch_run(model_cls: type[Model], 
              parameters: Mapping[str, Any | Iterable[Any]],
              number_processes: int | None = 1,
              iterations: int = 1,
              data_collection_period: int = -1,
              max_steps: int = 1000,
              display_progress: bool = True) -> list[dict[str, Any]]:
    """
    Batch run a Mesa model with parameter sweeps.
    
    Parameters:
        model_cls: The model class to run
        parameters: Dictionary of parameter names to values or ranges
        number_processes: Number of parallel processes (None for all cores)
        iterations: Number of iterations per parameter combination
        data_collection_period: How often to collect data (-1 for end only)
        max_steps: Maximum number of steps per model run
        display_progress: Whether to show progress bar
    
    Returns:
        List of dictionaries containing collected data
    """
    ...

from mesa import experimental

# Submodules available
experimental.devs          # Discrete Event Simulation framework
experimental.mesa_signals  # Reactive programming with observables
experimental.meta_agents   # Meta-agent functionality  
experimental.continuous_space  # Enhanced continuous space modeling

# Core type aliases used throughout Mesa
from typing import Any, Callable, Iterable, Sequence
import numpy as np
from random import Random

# Model-related types
SeedLike = int | np.integer | Sequence[int] | np.random.SeedSequence
RNGLike = np.random.Generator | np.random.BitGenerator

# Spatial types  
Coordinate = tuple[int, int]
FloatCoordinate = tuple[float, float] | np.ndarray
NetworkCoordinate = int
Position = Coordinate | FloatCoordinate | NetworkCoordinate

# Grid content types
GridContent = Agent | None
MultiGridContent = list[Agent]