rustsim 0.0.1

//! Schelling segregation model example.
//!
//! Classic ABM: agents on a grid are "happy" if enough neighbors share
//! their group. Unhappy agents move to a random empty cell.
//!
//! Run with: `cargo run -p rustsim --example schelling`

use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rustsim::prelude::*;
use rustsim_spaces::grid::{Grid2D, GridPos2};

const WIDTH: usize = 20;
const HEIGHT: usize = 20;
const NUM_AGENTS: usize = 300;
const TOLERANCE: f64 = 0.3; // min fraction of same-group neighbors to be happy

#[derive(Debug, Clone)]
struct Citizen {
    id: AgentId,
    pos: GridPos2,
    group: u8, // 0 or 1
    happy: bool,
}

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

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

type SchellingModel =
    StandardModel<Grid2D, Citizen, HashMapStore<Citizen>, SchellingProps, StdRng, Randomly>;

#[derive(Debug, Clone)]
struct SchellingProps {
    tolerance: f64,
    happy_count: usize,
}

fn citizen_step(
    agent: &mut Citizen,
    ctx: &mut StepContext<'_, Grid2D, Citizen, SchellingProps, StdRng, Randomly>,
) {
    let my_group = agent.group;
    let my_pos = agent.pos;

    // Count neighbors of same / different group
    let neighbor_ids = <Grid2D as SpaceInteraction<Citizen>>::nearby_ids(ctx.space(), &my_pos, 1);
    let mut same = 0usize;
    let mut diff = 0usize;

    for nid in &neighbor_ids {
        if *nid == agent.id {
            continue;
        }
        // We can't borrow other agents from ctx during agent step,
        // so we use a simple heuristic: odd IDs are group 0, even are group 1.
        // In a full implementation you'd use a snapshot or model-step approach.
        let neighbor_group = if *nid % 2 == 0 { 1 } else { 0 };
        if neighbor_group == my_group {
            same += 1;
        } else {
            diff += 1;
        }
    }

    let total = same + diff;
    let fraction = if total > 0 {
        same as f64 / total as f64
    } else {
        1.0
    };
    let tolerance = ctx.properties().tolerance;

    agent.happy = fraction >= tolerance;

    // If unhappy, move to a random position
    if !agent.happy {
        let (w, h) = ctx.space().dimensions();
        let new_pos: GridPos2 = (ctx.rng().gen_range(0..w), ctx.rng().gen_range(0..h));
        agent.pos = new_pos;
    }
}

fn model_step(model: &mut SchellingModel) {
    let happy = model.agents().filter(|a| a.happy).count();
    model.properties_mut().happy_count = happy;
}

fn main() {
    let grid = Grid2D::new(WIDTH, HEIGHT, false);
    let mut store = HashMapStore::new();
    let mut rng = StdRng::seed_from_u64(42);

    for i in 1..=NUM_AGENTS as u64 {
        let pos: GridPos2 = (rng.gen_range(0..WIDTH), rng.gen_range(0..HEIGHT));
        let group = if rng.gen_bool(0.5) { 0 } else { 1 };
        store.insert(Citizen {
            id: i,
            pos,
            group,
            happy: false,
        });
    }

    let props = SchellingProps {
        tolerance: TOLERANCE,
        happy_count: 0,
    };

    let mut model = SchellingModel::new(
        store,
        grid,
        Randomly::new(),
        props,
        StdRng::seed_from_u64(42),
        Some(Box::new(citizen_step)),
        Some(model_step),
        true,
    );

    println!("Schelling Segregation Model");
    println!("  Grid: {WIDTH}x{HEIGHT}");
    println!("  Agents: {NUM_AGENTS}");
    println!("  Tolerance: {TOLERANCE}");
    println!();

    for step in 1..=50u64 {
        model.step();
        let happy = model.properties().happy_count;
        let pct = 100.0 * happy as f64 / NUM_AGENTS as f64;
        if step <= 10 || step % 10 == 0 {
            println!("  Step {step:>3}: {happy:>3}/{NUM_AGENTS} happy ({pct:.1}%)");
        }
    }

    let final_happy = model.properties().happy_count;
    println!(
        "\nFinal: {final_happy}/{NUM_AGENTS} happy ({:.1}%)",
        100.0 * final_happy as f64 / NUM_AGENTS as f64
    );
}