rustsim 0.0.1

use rustsim::prelude::*;
use rustsim_spaces::graph::{GraphPos, GraphSpace, NeighborType};
use std::collections::HashMap;
use std::time::Instant;

#[derive(Debug, Clone)]
struct Particle {
    id: AgentId, // AgentId is a type alias for u64. This AgentId comes from the rustsim prelude and is used to uniquely identify agents in the model.
    x: f32,
    vx: f32,
}

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

impl SoaExtractable for Particle {
    fn num_columns() -> usize {
        2
    }

    fn column_names() -> Vec<&'static str> {
        vec!["x", "vx"]
    }

    fn extract_row(&self, columns: &mut [Vec<f32>]) {
        columns[0].push(self.x);
        columns[1].push(self.vx);
    }

    fn write_back_row(&mut self, columns: &[&[f32]], row: usize) {
        self.x = columns[0][row];
        self.vx = columns[1][row];
    }
}

fn integrate_cpu(columns: &mut [Vec<f32>], n: usize) {
    let (x_col, rest) = columns.split_at_mut(1);
    let x = &mut x_col[0];
    let vx = &rest[0];
    for i in 0..n {
        x[i] += vx[i];
    }
}

#[derive(Debug, Clone)]
struct RoutedWalker {
    id: AgentId,
    pos: GraphPos,
    destination: GraphPos,
    arrived: bool,
}

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

impl PositionedAgent for RoutedWalker {
    type Position = GraphPos;

    fn position(&self) -> &Self::Position {
        &self.pos
    }

    fn set_position(&mut self, position: Self::Position) {
        self.pos = position;
    }
}

#[derive(Debug, Clone)]
struct RoutingProps {
    next_hop_to_dest: HashMap<GraphPos, GraphPos>,
    arrivals: usize,
}

type RoutedGraphModel = StandardModel<
    GraphSpace,
    RoutedWalker,
    VecStore<RoutedWalker>,
    RoutingProps,
    rand::rngs::StdRng,
    ById,
>;

fn routed_graph_step(
    agent: &mut RoutedWalker,
    ctx: &mut StepContext<'_, GraphSpace, RoutedWalker, RoutingProps, rand::rngs::StdRng, ById>,
) {
    if agent.arrived {
        return;
    }

    if agent.pos == agent.destination {
        agent.arrived = true;
        ctx.properties_mut().arrivals += 1;
        return;
    }

    if let Some(&next) = ctx.properties().next_hop_to_dest.get(&agent.pos) {
        if next != agent.pos {
            let dummy = RoutedWalker {
                id: agent.id,
                pos: agent.pos,
                destination: agent.destination,
                arrived: agent.arrived,
            };
            ctx.space_mut().remove_agent(&dummy).unwrap();
            agent.pos = next;
            ctx.space_mut().add_agent(agent).unwrap();
        }
    }

    if agent.pos == agent.destination && !agent.arrived {
        agent.arrived = true;
        ctx.properties_mut().arrivals += 1;
    }
}

fn build_small_routing_graph() -> GraphSpace {
    let mut graph = GraphSpace::new(8);
    for i in 0..7 {
        graph.add_edge(i, i + 1);
    }
    graph
}

fn compute_next_hop_map(graph: &GraphSpace, destination: GraphPos) -> HashMap<GraphPos, GraphPos> {
    let mut map = HashMap::new();

    for start in 0..graph.num_vertices() {
        if start == destination {
            map.insert(start, start);
            continue;
        }

        let result = astar(
            start,
            destination,
            |a: &usize, b: &usize| (*a).abs_diff(*b) as f64,
            |node: &usize| {
                graph
                    .neighbors(*node, NeighborType::Out)
                    .into_iter()
                    .map(|n| (n, 1.0))
                    .collect::<Vec<_>>()
            },
        )
        .expect("path to destination should exist on the small routing graph");

        let next = *result
            .path
            .get(1)
            .expect("non-destination nodes should have a next hop");
        map.insert(start, next);
    }

    map
}

#[test]
fn cpu_batch_step_soak_preserves_expected_values() {
    const AGENT_COUNT: u64 = 25_000;
    const STEP_COUNT: usize = 120;

    let mut store = VecStore::new();
    for id in 1..=AGENT_COUNT {
        store.insert(Particle {
            id,
            x: 0.0,
            vx: 0.005 + (id % 5) as f32 * 0.001,
        });
    }

    let t0 = Instant::now();
    for _ in 0..STEP_COUNT {
        cpu_batch_step::<Particle, _, _>(&store, integrate_cpu);
    }
    let elapsed_ms = t0.elapsed().as_millis();

    let a1 = store.get(1).unwrap();
    let expected = STEP_COUNT as f32 * 0.006;
    assert!((a1.x - expected).abs() < 1e-4);
    drop(a1);

    let checksum: f64 = (1..=AGENT_COUNT)
        .map(|id| store.get(id).unwrap().x as f64)
        .sum();
    assert!(checksum.is_finite());
    assert!(checksum > 0.0);

    eprintln!(
        "[batch soak] {} agents x {} steps completed in {} ms",
        AGENT_COUNT, STEP_COUNT, elapsed_ms
    );
}

#[test]
fn device_soa_store_soak_preserves_state_across_many_cpu_steps() {
    const AGENT_COUNT: u64 = 50_000;
    const STEP_COUNT: usize = 150;

    let mut store = VecStore::new();
    for id in 1..=AGENT_COUNT {
        store.insert(Particle {
            id,
            x: id as f32 * 0.01,
            vx: 0.01,
        });
    }

    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    assert_eq!(device.agent_count(), AGENT_COUNT as usize);

    let t0 = Instant::now();
    for _ in 0..STEP_COUNT {
        let kernel_us = device.step_cpu(integrate_cpu);
        assert!(kernel_us > 0 || device.agent_count() > 0);
    }
    let elapsed_ms = t0.elapsed().as_millis();

    assert!(device.is_dirty());
    assert_eq!(device.agent_count(), AGENT_COUNT as usize);

    device.download::<Particle, _>(&store);
    assert!(!device.is_dirty());

    let a1 = store.get(1).unwrap();
    let expected = 0.01 + STEP_COUNT as f32 * 0.01;
    assert!((a1.x - expected).abs() < 1e-4);
    drop(a1);

    eprintln!(
        "[device soak] {} agents x {} steps completed in {} ms",
        AGENT_COUNT, STEP_COUNT, elapsed_ms
    );
}

#[test]
#[ignore]
fn scale_profile_device_soa_large_population() {
    const AGENT_COUNT: u64 = 300_000;
    const STEP_COUNT: usize = 60;

    let mut store = VecStore::new();
    for id in 1..=AGENT_COUNT {
        store.insert(Particle {
            id,
            x: 0.0,
            vx: 0.001,
        });
    }

    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);

    let t0 = Instant::now();
    for _ in 0..STEP_COUNT {
        device.step_cpu(integrate_cpu);
    }
    let elapsed_ms = t0.elapsed().as_millis();

    device.download::<Particle, _>(&store);
    assert!((store.get(1).unwrap().x - (STEP_COUNT as f32 * 0.001)).abs() < 1e-5);

    eprintln!(
        "[device scale] {} agents x {} steps completed in {} ms",
        AGENT_COUNT, STEP_COUNT, elapsed_ms
    );
}

#[test]
#[ignore]
fn scale_profile_one_million_agents_small_graph_pathfinding_to_destination() {
    use rand::SeedableRng;

    const AGENT_COUNT: u64 = 1_000_000;
    const DESTINATION: GraphPos = 7;
    const STEPS_TO_DESTINATION: usize = 7;

    let space = build_small_routing_graph();
    let next_hop_to_dest = compute_next_hop_map(&space, DESTINATION);

    let mut store = VecStore::new();
    for id in 1..=AGENT_COUNT {
        let start = ((id - 1) % DESTINATION as u64) as usize;
        store.insert(RoutedWalker {
            id,
            pos: start,
            destination: DESTINATION,
            arrived: false,
        });
    }

    let props = RoutingProps {
        next_hop_to_dest,
        arrivals: 0,
    };

    let mut model = RoutedGraphModel::new_with_agent_step(
        store,
        space,
        ById::new(),
        props,
        rand::rngs::StdRng::seed_from_u64(42),
        routed_graph_step,
        true,
    );

    let t0 = Instant::now();
    model.step_n(STEPS_TO_DESTINATION);
    let elapsed_ms = t0.elapsed().as_millis();

    assert_eq!(model.agents_len(), AGENT_COUNT as usize);
    assert_eq!(model.properties().arrivals, AGENT_COUNT as usize);

    let arrived = model.agents().filter(|a| a.arrived).count();
    let at_destination = model.agents().filter(|a| a.pos == DESTINATION).count();
    assert_eq!(arrived, AGENT_COUNT as usize);
    assert_eq!(at_destination, AGENT_COUNT as usize);

    let destination_occupancy = model.space().ids_in_position(DESTINATION).len();
    assert_eq!(destination_occupancy, AGENT_COUNT as usize);

    let checksum: u64 = model
        .agents()
        .map(|a| a.pos as u64 + if a.arrived { 1 } else { 0 })
        .sum();
    assert_eq!(checksum, AGENT_COUNT * (DESTINATION as u64 + 1));

    eprintln!(
        "[graph pathfinding scale] {} agents reached node {} on small graph in {} steps / {} ms",
        AGENT_COUNT, DESTINATION, STEPS_TO_DESTINATION, elapsed_ms
    );
}