rdpe 0.1.0

//! # CPU Readback Demo
//!
//! Demonstrates reading particle data back from GPU to CPU for aggregate
//! computations that would be complex or impossible in shaders alone.
//!
//! ## What This Demonstrates
//!
//! - `ctx.request_readback()` - Request GPU particle data for next frame
//! - `ctx.particles_raw()` - Access raw bytes of particle data
//! - `bytemuck::cast_slice()` - Convert bytes to typed particle slice
//! - CPU-computed aggregates fed back as uniforms
//!
//! ## The Simulation
//!
//! Two particle types:
//! - **Leaders (green)**: Wander randomly through space
//! - **Followers (blue)**: Chase the center of mass of all leaders
//!
//! The center of mass is computed on CPU from readback data and fed back
//! to the GPU as a uniform. This creates a feedback loop:
//! GPU simulation -> CPU analysis -> GPU uniforms -> GPU simulation
//!
//! ## Performance Note
//!
//! Readback stalls the GPU pipeline. This example requests readback every
//! 10 frames (~6 times/second at 60fps). For production, consider:
//! - Lower frequency (once per second)
//! - On-demand (keypress triggered)
//! - Async readback (more complex, no stall)
//!
//! Run with: `cargo run --example cpu_readback`

use rand::Rng;
use rdpe::prelude::*;

#[derive(ParticleType, Clone, Copy, PartialEq)]
enum Role {
    Leader,   // 0 - Wanders randomly, others follow
    Follower, // 1 - Chases center of mass of leaders
}

#[derive(Particle, Clone)]
struct Agent {
    position: Vec3,
    velocity: Vec3,
    #[color]
    color: Vec3,
    particle_type: u32,
}

/// GPU-compatible struct layout (generated by #[derive(Particle)])
/// Must match exactly for bytemuck::cast_slice to work correctly.
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct AgentGpuLayout {
    position: [f32; 3],
    _pad0: f32,
    velocity: [f32; 3],
    _pad1: f32,
    color: [f32; 3],
    _pad2: f32,
    particle_type: u32,
    age: f32,
    alive: u32,
    scale: f32,
}

fn main() {
    let mut rng = rand::thread_rng();

    let num_leaders = 500;
    let num_followers = 2000;
    let total = num_leaders + num_followers;

    // Create particles
    let particles: Vec<Agent> = (0..total)
        .map(|i| {
            let is_leader = i < num_leaders;
            let role = if is_leader {
                Role::Leader
            } else {
                Role::Follower
            };

            let pos = Vec3::new(
                rng.gen_range(-0.8..0.8),
                rng.gen_range(-0.8..0.8),
                rng.gen_range(-0.8..0.8),
            );

            let vel = Vec3::new(
                rng.gen_range(-0.2..0.2),
                rng.gen_range(-0.2..0.2),
                rng.gen_range(-0.2..0.2),
            );

            let color = if is_leader {
                Vec3::new(0.2, 0.9, 0.3) // Green leaders
            } else {
                Vec3::new(0.3, 0.5, 1.0) // Blue followers
            };

            Agent {
                position: pos,
                velocity: vel,
                color,
                particle_type: role.into(),
            }
        })
        .collect();

    // Track frame count for periodic readback
    let mut frame_count = 0u32;

    Simulation::<Agent>::new()
        .with_particle_count(total as u32)
        .with_bounds(1.0)
        .with_spatial_config(0.15, 32)
        .with_spawner(move |ctx| particles[ctx.index as usize].clone())
        // Uniforms for center of mass (computed on CPU)
        .with_uniform("leader_com_x", 0.0f32)
        .with_uniform("leader_com_y", 0.0f32)
        .with_uniform("leader_com_z", 0.0f32)
        .with_uniform("leader_count", num_leaders as f32)
        // Update callback - CPU analysis happens here
        .with_update(move |ctx| {
            frame_count += 1;

            // Request readback every 10 frames
            if frame_count % 10 == 0 {
                ctx.request_readback();
            }

            // Process readback data using with_particles() for clean borrow handling
            // The closure computes stats and returns them, then we can call ctx.set()
            if let Some((com, count, avg_vel)) = ctx.with_particles(|bytes| {
                // Cast raw bytes to our GPU struct layout
                let particles: &[AgentGpuLayout] = rdpe::bytemuck::cast_slice(bytes);

                // Compute center of mass of leaders
                let mut sum = Vec3::ZERO;
                let mut count = 0;

                for p in particles.iter() {
                    // particle_type 0 = Leader
                    if p.particle_type == 0 && p.alive == 1 {
                        sum += Vec3::from_array(p.position);
                        count += 1;
                    }
                }

                // Compute average velocity (for logging)
                let avg_vel: f32 = particles
                    .iter()
                    .filter(|p| p.alive == 1)
                    .map(|p| {
                        let v = Vec3::from_array(p.velocity);
                        v.length()
                    })
                    .sum::<f32>()
                    / particles.len().max(1) as f32;

                let com = if count > 0 {
                    sum / count as f32
                } else {
                    Vec3::ZERO
                };

                (com, count, avg_vel)
            }) {
                if count > 0 {
                    ctx.set("leader_com_x", com.x);
                    ctx.set("leader_com_y", com.y);
                    ctx.set("leader_com_z", com.z);
                    ctx.set("leader_count", count as f32);

                    // Print stats occasionally
                    if frame_count % 60 == 0 {
                        println!(
                            "Leaders: {} | CoM: ({:.2}, {:.2}, {:.2}) | Avg speed: {:.3}",
                            count, com.x, com.y, com.z, avg_vel
                        );
                    }
                }
            }
        })
        // Leaders wander randomly
        .with_rule(Rule::Typed {
            self_type: Role::Leader.into(),
            other_type: None,
            rule: Box::new(Rule::Wander {
                strength: 2.0,
                frequency: 50.0,
            }),
        })
        // Followers chase the center of mass (uniform-based)
        .with_rule(Rule::Custom(
            r#"
            // Only apply to followers (type 1)
            if p.particle_type == 1u {
                let com = vec3<f32>(
                    uniforms.leader_com_x,
                    uniforms.leader_com_y,
                    uniforms.leader_com_z
                );

                let to_com = com - p.position;
                let dist = length(to_com);

                if dist > 0.01 {
                    // Chase the center of mass
                    let dir = normalize(to_com);
                    p.velocity += dir * 3.0 * uniforms.delta_time;
                }
            }
            "#
            .to_string(),
        ))
        // Everyone maintains some separation
        .with_rule(Rule::Separate {
            radius: 0.08,
            strength: 2.0,
        })
        // Physics
        .with_rule(Rule::SpeedLimit { min: 0.0, max: 1.5 })
        .with_rule(Rule::Drag(1.5))
        .with_rule(Rule::BounceWalls { restitution: 1.0 })
        .run().expect("Simulation failed");
}