rustsim-spaces 0.0.1

use rand::rngs::StdRng;
use rand::SeedableRng;
use rustsim_core::interaction::{PositionedAgent, SpaceInteraction};
use rustsim_core::prelude::*;
use rustsim_spaces::graph::{GraphPos, GraphSpace, GraphSpaceError, NeighborType};

// ---- Agent for graph space ----

#[derive(Debug, Clone)]
struct GraphBot {
    id: AgentId,
    pos: GraphPos,
}

impl Agent for GraphBot {
    fn id(&self) -> AgentId {
        self.id
    }
}

impl PositionedAgent for GraphBot {
    type Position = GraphPos;
    fn position(&self) -> &GraphPos {
        &self.pos
    }
    fn set_position(&mut self, p: GraphPos) {
        self.pos = p;
    }
}

// ====================================================================
// Basic graph operations
// ====================================================================

#[test]
fn create_graph_add_edges() {
    let mut g = GraphSpace::new(5);
    assert_eq!(g.num_vertices(), 5);
    assert_eq!(g.num_edges(), 0);

    g.add_edge(0, 1);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);

    assert_eq!(g.num_edges(), 4);
    assert_eq!(g.neighbors_out(1), &[0, 2]); // undirected: both directions
}

#[test]
fn directed_graph() {
    let mut g = GraphSpace::new_directed(3, true);
    g.add_edge(0, 1);
    g.add_edge(1, 2);

    assert_eq!(g.num_edges(), 2);
    assert_eq!(g.neighbors_out(0), &[1]);
    assert!(g.neighbors_out(1).contains(&2));
    assert_eq!(g.neighbors_in(2), &[1]);
    // Node 0 has no in-neighbors
    assert!(g.neighbors_in(0).is_empty());
}

#[test]
fn add_remove_vertex() {
    let mut g = GraphSpace::new(3);
    g.add_edge(0, 1);
    g.add_edge(1, 2);

    let new_v = g.add_vertex();
    assert_eq!(new_v, 3);
    assert_eq!(g.num_vertices(), 4);

    g.add_edge(2, 3);
    assert_eq!(g.num_edges(), 3);

    // Remove vertex 1 (swaps with last vertex 3)
    g.rem_vertex(1);
    assert_eq!(g.num_vertices(), 3);
    // Edge 0-1 and 1-2 were removed; edge 2-3 was remapped
}

#[test]
fn add_remove_edge() {
    let mut g = GraphSpace::new(4);
    g.add_edge(0, 1);
    g.add_edge(1, 2);
    g.add_edge(2, 3);

    assert_eq!(g.num_edges(), 3);
    g.rem_edge(1, 2);
    assert_eq!(g.num_edges(), 2);

    // Duplicate add returns false
    assert!(!g.add_edge(0, 1));
}

// ====================================================================
// Agent placement and neighbor queries
// ====================================================================

#[test]
fn agents_on_graph() {
    let mut g = GraphSpace::new(4);
    g.add_edge(0, 1);
    g.add_edge(1, 2);
    g.add_edge(2, 3);

    let a = GraphBot { id: 1, pos: 0 };
    let b = GraphBot { id: 2, pos: 1 };
    let c = GraphBot { id: 3, pos: 3 };

    SpaceInteraction::<GraphBot>::add_agent(&mut g, &a).unwrap();
    SpaceInteraction::<GraphBot>::add_agent(&mut g, &b).unwrap();
    SpaceInteraction::<GraphBot>::add_agent(&mut g, &c).unwrap();

    assert_eq!(g.ids_in_position(0), &[1]);
    assert_eq!(g.ids_in_position(1), &[2]);
    assert_eq!(g.ids_in_position(3), &[3]);

    // Nearby ids at node 1, radius 1: includes node 1 itself + neighbors 0, 2
    let near = g.nearby_agent_ids(1, 1, NeighborType::Out);
    assert!(near.contains(&1)); // agent at node 0
    assert!(near.contains(&2)); // agent at node 1
                                // node 2 is neighbor of 1, but no agent there
    assert!(!near.contains(&3)); // agent at node 3 is 2 hops away

    // Radius 2: now includes node 3
    let near2 = g.nearby_agent_ids(1, 2, NeighborType::Out);
    assert!(near2.contains(&3));
}

#[test]
fn remove_agent_from_graph() {
    let mut g = GraphSpace::new(2);
    g.add_edge(0, 1);

    let a = GraphBot { id: 1, pos: 0 };
    SpaceInteraction::<GraphBot>::add_agent(&mut g, &a).unwrap();
    assert_eq!(g.ids_in_position(0).len(), 1);

    SpaceInteraction::<GraphBot>::remove_agent(&mut g, &a).unwrap();
    assert!(g.ids_in_position(0).is_empty());
}

#[test]
fn nearby_positions_bfs() {
    // Linear chain: 0 - 1 - 2 - 3 - 4
    let mut g = GraphSpace::new(5);
    for i in 0..4 {
        g.add_edge(i, i + 1);
    }

    let r1 = g.nearby_positions(2, 1, NeighborType::Out);
    assert!(r1.contains(&1));
    assert!(r1.contains(&3));
    assert!(!r1.contains(&0));
    assert!(!r1.contains(&4));
    assert!(!r1.contains(&2)); // origin excluded

    let r2 = g.nearby_positions(2, 2, NeighborType::Out);
    assert!(r2.contains(&0));
    assert!(r2.contains(&4));
}

// ====================================================================
// Multi-floor building scenario
// ====================================================================

#[test]
fn multi_floor_building() {
    // Model a 2-floor building:
    //
    //   Floor 2:  [Room2A=4] -- [Corridor2=5] -- [Room2B=6]
    //                                |
    //                            [Stairs=3]
    //                                |
    //   Floor 1:  [Room1A=0] -- [Corridor1=1] -- [Room1B=2]
    //
    let mut g = GraphSpace::new(7);

    // Floor 1 edges
    g.add_edge(0, 1); // Room1A <-> Corridor1
    g.add_edge(1, 2); // Corridor1 <-> Room1B

    // Stairs connecting floors
    g.add_edge(1, 3); // Corridor1 <-> Stairs
    g.add_edge(3, 5); // Stairs <-> Corridor2

    // Floor 2 edges
    g.add_edge(4, 5); // Room2A <-> Corridor2
    g.add_edge(5, 6); // Corridor2 <-> Room2B

    assert_eq!(g.num_vertices(), 7);
    assert_eq!(g.num_edges(), 6);

    // Place agents
    let agents = vec![
        GraphBot { id: 1, pos: 0 }, // Room1A
        GraphBot { id: 2, pos: 2 }, // Room1B
        GraphBot { id: 3, pos: 4 }, // Room2A
        GraphBot { id: 4, pos: 6 }, // Room2B
    ];

    for a in &agents {
        SpaceInteraction::<GraphBot>::add_agent(&mut g, a).unwrap();
    }

    // From Corridor1 (node 1), radius 1: reach Room1A(0), Room1B(2), Stairs(3)
    let near_corridor1 = g.nearby_positions(1, 1, NeighborType::Out);
    assert!(near_corridor1.contains(&0));
    assert!(near_corridor1.contains(&2));
    assert!(near_corridor1.contains(&3));
    assert!(!near_corridor1.contains(&5)); // Corridor2 is 2 hops away

    // From Corridor1, radius 2: also reach Corridor2(5) via stairs
    let near_corridor1_r2 = g.nearby_positions(1, 2, NeighborType::Out);
    assert!(near_corridor1_r2.contains(&5));

    // Agent 1 (Room1A) can "see" agent 2 (Room1B) within 2 hops
    let near_room1a = g.nearby_agent_ids(0, 2, NeighborType::Out);
    assert!(near_room1a.contains(&1)); // self at node 0
    assert!(near_room1a.contains(&2)); // agent at node 2

    // Agent 1 cannot see agents on Floor 2 within 2 hops
    // Room1A -> Corridor1 -> Stairs = 2 hops, Stairs has no agents
    assert!(!near_room1a.contains(&3));
    assert!(!near_room1a.contains(&4));

    // But within 4 hops (Room1A -> Corr1 -> Stairs -> Corr2 -> Room2A)
    let near_room1a_r4 = g.nearby_agent_ids(0, 4, NeighborType::Out);
    assert!(near_room1a_r4.contains(&3)); // agent at Room2A (node 4)
    assert!(near_room1a_r4.contains(&4)); // agent at Room2B (node 6)
}

// ====================================================================
// Random position and SpaceInteraction trait
// ====================================================================

#[test]
fn random_position_in_graph() {
    let g = GraphSpace::new(10);
    let mut rng = StdRng::seed_from_u64(7);

    for _ in 0..100 {
        let pos = SpaceInteraction::<GraphBot>::random_position(&g, &mut rng);
        assert!(pos < 10);
    }
}

#[test]
fn invalid_node_returns_error() {
    let mut g = GraphSpace::new(3);
    let a = GraphBot { id: 1, pos: 99 };
    let result = SpaceInteraction::<GraphBot>::add_agent(&mut g, &a);
    assert!(matches!(result, Err(GraphSpaceError::InvalidNode(99))));
}