from mesa import Agent, Model
from mesa.experimental.continuous_space import ContinuousSpace
from mesa.experimental.continuous_space.continuous_space_agents import ContinuousSpaceAgent
import numpy as np

class DummyAgent(ContinuousSpaceAgent):
    pass

model = Model()
space = ContinuousSpace(dimensions=[[0, 10], [0, 10]], torus=False)

# Create 3 agents
agents = []
for i in range(3):
    agent = DummyAgent(space, model)
    agent.position = np.array([float(i), float(i)])
    agents.append(agent)

print(f"Initial agents: {len(space.active_agents)}")
print(f"Initial index map size: {len(space._index_to_agent)}")

# Remove the middle agent (index 1)
agent_to_remove = agents[1]
print(f"Removing agent at index: {space._agent_to_index[agent_to_remove]}")
space._remove_agent(agent_to_remove) 

print(f"Agents after removal: {len(space.active_agents)}")
print(f"Index map size after removal: {len(space._index_to_agent)}")

# Check for stale index
expected_size = len(space.active_agents)
actual_size = len(space._index_to_agent)

if actual_size > expected_size:
    print("BUG DETECTED: _index_to_agent has stale entries!")
    print(f"Expected size: {expected_size}, Actual size: {actual_size}")
    last_index = actual_size - 1
    print(f"Entry at index {last_index}: {space._index_to_agent.get(last_index)}")