rustsim 0.0.1

//! Social Force Model (Helbing & Molnar 1995) for pedestrian dynamics.
//!
//! This example demonstrates how to implement the classic Social Force Model
//! using rustsim's `ContinuousSpace2D` and spatial neighbor queries. Pedestrians:
//!
//! - Are driven toward a destination by a **desired force**
//! - Are repelled from other pedestrians by a **social repulsion force**
//! - Are repelled from walls by a **wall repulsion force**
//!
//! ## References
//!
//! - Helbing, D. & Molnar, P. (1995). "Social force model for pedestrian
//!   dynamics." Physical Review E, 51(5), 4282.
//! - Helbing, D., Farkas, I., & Vicsek, T. (2000). "Simulating dynamical
//!   features of escape panic." Nature, 407(6803), 487-490.
//!
//! Run with: `cargo run -p rustsim --example social_force --release`

use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rustsim::prelude::*;
use rustsim_spaces::continuous::{ContinuousPos, ContinuousSpace2D};

// ---- Simulation parameters ----

const NUM_AGENTS: usize = 200;
const SPACE_X: f64 = 40.0; // meters
const SPACE_Y: f64 = 20.0; // meters
const DT: f64 = 0.1; // time step in seconds
const NUM_STEPS: usize = 2000; // 200 seconds total
const SEARCH_RADIUS: f64 = 5.0; // meters

// Social Force Model parameters (Helbing & Molnar 1995)
const TAU: f64 = 0.5; // relaxation time (seconds)
const DESIRED_SPEED: f64 = 1.3; // desired walking speed (m/s)
const A_SOC: f64 = 2.1; // social force strength (N)
const B_SOC: f64 = 0.3; // social force range (m)
const A_WALL: f64 = 10.0; // wall force strength (N)
const B_WALL: f64 = 0.2; // wall force range (m)
const AGENT_RADIUS: f64 = 0.25; // pedestrian body radius (m)
const MAX_SPEED: f64 = 2.5; // maximum speed (m/s)

// ---- Agent ----

#[derive(Debug, Clone)]
struct Pedestrian {
    id: AgentId,
    x: f64,
    y: f64,
    vx: f64,
    vy: f64,
    dest_x: f64,
    dest_y: f64,
    desired_speed: f64,
    radius: f64,
    arrived: bool,
}

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

// ---- Wall definition ----

/// A wall segment defined by two endpoints.
#[derive(Debug, Clone)]
struct Wall {
    x1: f64,
    y1: f64,
    x2: f64,
    y2: f64,
}

impl Wall {
    fn new(x1: f64, y1: f64, x2: f64, y2: f64) -> Self {
        Self { x1, y1, x2, y2 }
    }

    /// Closest point on this wall segment to point (px, py).
    fn closest_point(&self, px: f64, py: f64) -> (f64, f64) {
        let dx = self.x2 - self.x1;
        let dy = self.y2 - self.y1;
        let len_sq = dx * dx + dy * dy;
        if len_sq < 1e-12 {
            return (self.x1, self.y1);
        }
        let t = ((px - self.x1) * dx + (py - self.y1) * dy) / len_sq;
        let t = t.clamp(0.0, 1.0);
        (self.x1 + t * dx, self.y1 + t * dy)
    }
}

// ---- Model properties ----

#[derive(Debug, Clone)]
struct SfmProps {
    walls: Vec<Wall>,
    dt: f64,
    total_arrived: usize,
    density_grid: DensityGrid,
}

// ---- Model type ----

type SfmModel = StandardModel<
    ContinuousSpace2D,
    Pedestrian,
    HashMapStore<Pedestrian>,
    SfmProps,
    StdRng,
    Fastest,
>;

// ---- Agent step: Social Force Model ----

fn sfm_step(
    agent: &mut Pedestrian,
    ctx: &mut StepContext<'_, ContinuousSpace2D, Pedestrian, SfmProps, StdRng, Fastest>,
) {
    if agent.arrived {
        return;
    }

    let dt = ctx.properties().dt;

    // 1. DESIRED FORCE: f_desired = (v_desired * e_i - v_i) / tau
    let dx_dest = agent.dest_x - agent.x;
    let dy_dest = agent.dest_y - agent.y;
    let dist_dest = (dx_dest * dx_dest + dy_dest * dy_dest).sqrt();

    if dist_dest < 0.5 {
        agent.arrived = true;
        agent.vx = 0.0;
        agent.vy = 0.0;
        return;
    }

    let ex = dx_dest / dist_dest;
    let ey = dy_dest / dist_dest;

    let f_desired_x = (agent.desired_speed * ex - agent.vx) / TAU;
    let f_desired_y = (agent.desired_speed * ey - agent.vy) / TAU;

    // 2. SOCIAL REPULSION FORCE from nearby pedestrians
    // Use space's neighbor query for position-based repulsion.
    let nearby = ctx
        .space()
        .nearby_ids_euclidean(&ContinuousPos::new(agent.x, agent.y), SEARCH_RADIUS);

    let mut f_social_x = 0.0;
    let mut f_social_y = 0.0;

    for &nid in &nearby {
        if nid == agent.id {
            continue;
        }
        if let Some(npos) = ctx.space().agent_position(nid) {
            let dx = agent.x - npos.x;
            let dy = agent.y - npos.y;
            let dist = (dx * dx + dy * dy).sqrt();
            let r_sum = agent.radius + AGENT_RADIUS;
            if dist > 0.01 {
                let nx = dx / dist;
                let ny = dy / dist;
                // Exponential repulsion: A * exp(-(d - r_sum) / B)
                let force = A_SOC * (-(dist - r_sum) / B_SOC).exp();
                f_social_x += force * nx;
                f_social_y += force * ny;
            }
        }
    }

    // 3. WALL REPULSION FORCE
    let mut f_wall_x = 0.0;
    let mut f_wall_y = 0.0;

    for wall in &ctx.properties().walls {
        let (cx, cy) = wall.closest_point(agent.x, agent.y);
        let dx = agent.x - cx;
        let dy = agent.y - cy;
        let dist = (dx * dx + dy * dy).sqrt();

        if dist > 0.01 && dist < 3.0 {
            let nx = dx / dist;
            let ny = dy / dist;
            // Exponential wall repulsion
            let force = A_WALL * (-(dist - agent.radius) / B_WALL).exp();
            f_wall_x += force * nx;
            f_wall_y += force * ny;
        }
    }

    // 4. INTEGRATE: Euler integration
    let ax = f_desired_x + f_social_x + f_wall_x;
    let ay = f_desired_y + f_social_y + f_wall_y;

    agent.vx += ax * dt;
    agent.vy += ay * dt;

    // Clamp speed
    let speed = (agent.vx * agent.vx + agent.vy * agent.vy).sqrt();
    if speed > MAX_SPEED {
        let scale = MAX_SPEED / speed;
        agent.vx *= scale;
        agent.vy *= scale;
    }

    agent.x += agent.vx * dt;
    agent.y += agent.vy * dt;

    // Clamp to space bounds
    agent.x = agent.x.clamp(0.01, SPACE_X - 0.01);
    agent.y = agent.y.clamp(0.01, SPACE_Y - 0.01);
}

// ---- Model step: update density grid, remove arrived agents ----

fn model_step(model: &mut SfmModel) {
    // Update density grid
    model.properties_mut().density_grid.clear();
    let positions: Vec<(f64, f64)> = model
        .agents()
        .filter(|a| !a.arrived)
        .map(|a| (a.x, a.y))
        .collect();
    model.properties_mut().density_grid.add_positions(positions);

    // Remove arrived agents and count them
    let to_remove: Vec<AgentId> = model
        .agents()
        .filter(|a| a.arrived)
        .map(|a| a.id())
        .collect();
    let newly_arrived = to_remove.len();
    for id in to_remove {
        model.remove_agent(id);
    }
    model.properties_mut().total_arrived += newly_arrived;

    // Update space positions for remaining agents
    let agent_data: Vec<(AgentId, f64, f64)> = model.agents().map(|a| (a.id(), a.x, a.y)).collect();
    for (id, x, y) in agent_data {
        let _ = model
            .space_mut()
            .move_agent_pos(id, ContinuousPos::new(x, y));
    }
}

// ---- Main ----

fn main() {
    let mut rng = StdRng::seed_from_u64(42);

    // Create continuous space with spatial hashing
    let mut space = ContinuousSpace2D::new(SPACE_X, SPACE_Y, false, SEARCH_RADIUS).unwrap();

    // Define walls: corridor boundaries with a bottleneck in the middle
    let walls = vec![
        // Top boundary
        Wall::new(0.0, SPACE_Y, SPACE_X, SPACE_Y),
        // Bottom boundary
        Wall::new(0.0, 0.0, SPACE_X, 0.0),
        // Bottleneck: upper obstacle
        Wall::new(18.0, SPACE_Y, 18.0, SPACE_Y * 0.65),
        Wall::new(22.0, SPACE_Y, 22.0, SPACE_Y * 0.65),
        // Bottleneck: lower obstacle
        Wall::new(18.0, 0.0, 18.0, SPACE_Y * 0.35),
        Wall::new(22.0, 0.0, 22.0, SPACE_Y * 0.35),
    ];

    // Create agents - two counter-flowing streams
    let mut store = HashMapStore::new();
    for i in 1..=NUM_AGENTS as u64 {
        let going_right = i <= (NUM_AGENTS as u64) / 2;
        let (x, dest_x) = if going_right {
            (rng.gen_range(1.0..8.0), SPACE_X - 1.0)
        } else {
            (rng.gen_range(SPACE_X - 8.0..SPACE_X - 1.0), 1.0)
        };
        let y = rng.gen_range(SPACE_Y * 0.4..SPACE_Y * 0.6);

        let ped = Pedestrian {
            id: i,
            x,
            y,
            vx: 0.0,
            vy: 0.0,
            dest_x,
            dest_y: y + rng.gen_range(-1.0..1.0),
            desired_speed: DESIRED_SPEED + rng.gen_range(-0.2..0.2),
            radius: AGENT_RADIUS,
            arrived: false,
        };

        // Register with space using a helper wrapper
        let pos = ContinuousPos::new(ped.x, ped.y);
        let wrapper = PosWrapper { id: ped.id, pos };
        <ContinuousSpace2D as SpaceInteraction<PosWrapper>>::add_agent(&mut space, &wrapper)
            .expect("initial placement should succeed");

        store.insert(ped);
    }

    let props = SfmProps {
        walls,
        dt: DT,
        total_arrived: 0,
        density_grid: DensityGrid::new(SPACE_X, SPACE_Y, 2.0).unwrap(),
    };

    let mut model = SfmModel::new(
        store,
        space,
        Fastest::new(),
        props,
        StdRng::seed_from_u64(42),
        Some(Box::new(sfm_step)),
        Some(model_step),
        true, // agents step first, then model step updates space
    );

    // LoS criteria for pedestrian walkway analysis
    let walkway_los = PedestrianWalkway;

    // ---- Run simulation ----

    println!();
    println!("=== Social Force Model (Helbing & Molnar 1995) ===");
    println!(
        "Pedestrians: {} ({} per direction)",
        NUM_AGENTS,
        NUM_AGENTS / 2
    );
    println!(
        "Space:       {:.0} x {:.0} m with bottleneck at x=18-22",
        SPACE_X, SPACE_Y
    );
    println!("Time step:   {DT} s");
    println!(
        "Duration:    {} s ({NUM_STEPS} steps)",
        NUM_STEPS as f64 * DT
    );
    println!("Model:       desired force + social repulsion + wall repulsion");
    println!(
        "LoS:         {} ({})",
        walkway_los.name(),
        walkway_los.unit()
    );
    println!();

    println!(
        "{:>6}  {:>6}  {:>10}  {:>10}  {:>6}  {:>6}",
        "Time", "Active", "Max rho", "Mean rho", "LoS", "Done"
    );
    println!("{}", "-".repeat(58));

    let t0 = std::time::Instant::now();

    for step in 0..NUM_STEPS {
        model.step();

        if step % 50 == 0 || step == NUM_STEPS - 1 {
            let time = (step + 1) as f64 * DT;
            let active = model.agents_len();
            let stats = model.properties().density_grid.statistics(&walkway_los);
            let arrived = model.properties().total_arrived;

            println!(
                "{:5.1}s  {:>6}  {:>7.2}/m2  {:>7.2}/m2  {:>5}  {:>6}",
                time, active, stats.max_density, stats.mean_density, stats.worst_los, arrived,
            );
        }
    }

    let elapsed = t0.elapsed();
    let total_arrived = model.properties().total_arrived;

    println!();
    println!("=== Results ===");
    println!("Pedestrians arrived: {total_arrived}/{NUM_AGENTS}");
    println!("Still walking:       {}", model.agents_len());

    let stats = model.properties().density_grid.statistics(&walkway_los);
    println!("Final max density:   {:.2} pax/m2", stats.max_density);
    println!("Final worst LoS:     {}", stats.worst_los);
    println!(
        "LoS distribution:    A={} B={} C={} D={} E={} F={}",
        stats.los_distribution[0],
        stats.los_distribution[1],
        stats.los_distribution[2],
        stats.los_distribution[3],
        stats.los_distribution[4],
        stats.los_distribution[5],
    );
    println!("Wall time:           {:.1} ms", elapsed.as_millis());
    println!(
        "Steps/s:             {:.0}",
        NUM_STEPS as f64 / elapsed.as_secs_f64()
    );
}

/// Helper wrapper for initial space registration.
/// (We need a PositionedAgent to register with ContinuousSpace2D.)
#[derive(Debug, Clone)]
struct PosWrapper {
    id: AgentId,
    pos: ContinuousPos,
}

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

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