elegans 1.0.0

//! Real-time 3D visualizer for the elegans worm simulation.
//!
//! Run: `cargo run --example visualize --features visualize`
//!
//! Controls:
//!   Space        — Pause / resume
//!   Up / Down    — Speed 1x → 64x
//!   Left drag    — Pan camera
//!   Right drag   — Orbit camera
//!   Scroll       — Zoom (smooth)
//!   R            — Reset camera
//!   F            — Follow worm (camera tracks center of mass)
//!   V            — Start / stop recording (saves frames as PNG)
//!   P            — Save single screenshot

use elegans::body::SEGMENT_COUNT;
use elegans::phase::DevelopmentalPhase;
use elegans::sim::{WormConfig, WormSim};
use macroquad::prelude::*;

use std::fs;
use std::path::PathBuf;

// ── Camera ──────────────────────────────────────────────────────────

struct OrbitCamera {
    target: Vec3,
    distance: f32,
    target_distance: f32,
    yaw: f32,
    pitch: f32,
    follow: bool,
    /// Previous mouse position for manual delta tracking.
    prev_mouse: Vec2,
}

impl OrbitCamera {
    fn new() -> Self {
        Self {
            target: Vec3::ZERO,
            distance: 30.0,
            target_distance: 30.0,
            yaw: 0.3,
            pitch: 0.6,
            follow: true,
            prev_mouse: Vec2::ZERO,
        }
    }

    fn reset(&mut self) {
        self.distance = 30.0;
        self.target_distance = 30.0;
        self.yaw = 0.3;
        self.pitch = 0.6;
        self.follow = true;
    }

    fn update(&mut self, worm_center: Vec3) {
        let mouse_pos: Vec2 = mouse_position().into();
        let delta = mouse_pos - self.prev_mouse;
        self.prev_mouse = mouse_pos;

        // Right-drag: orbit (rotate view)
        if is_mouse_button_down(MouseButton::Right) {
            self.yaw -= delta.x * 0.005;
            self.pitch = (self.pitch - delta.y * 0.005).clamp(0.05, 1.5);
        }

        // Left-drag: pan (move target in screen-plane)
        if is_mouse_button_down(MouseButton::Left) {
            let cam_fwd = (self.target - self.position()).normalize();
            let cam_right = cam_fwd.cross(Vec3::Z).normalize();
            let cam_up = cam_right.cross(cam_fwd).normalize();
            let pan_speed = self.distance * 0.003;
            self.target += cam_right * (-delta.x * pan_speed) + cam_up * (delta.y * pan_speed);
            self.follow = false;
        }

        // Scroll zoom: multiplicative, smoothly interpolated
        let (_, scroll_y) = mouse_wheel();
        if scroll_y.abs() > 0.01 {
            let factor = 1.0 - scroll_y * 0.08;
            self.target_distance = (self.target_distance * factor).clamp(3.0, 150.0);
        }
        self.distance += (self.target_distance - self.distance) * 0.15;

        if self.follow {
            self.target += (worm_center - self.target) * 0.1;
        }
    }

    fn position(&self) -> Vec3 {
        let x = self.distance * self.pitch.sin() * self.yaw.cos();
        let y = self.distance * self.pitch.sin() * self.yaw.sin();
        let z = self.distance * self.pitch.cos();
        self.target + Vec3::new(x, y, z)
    }

    fn camera3d(&self) -> Camera3D {
        Camera3D {
            position: self.position(),
            target: self.target,
            up: Vec3::Z,
            ..Default::default()
        }
    }
}

// ── Recording ───────────────────────────────────────────────────────

struct Recorder {
    recording: bool,
    dir: PathBuf,
    frame_count: u32,
    /// How many recording sessions have been started.
    session: u32,
}

impl Recorder {
    fn new() -> Self {
        Self {
            recording: false,
            dir: PathBuf::new(),
            frame_count: 0,
            session: 0,
        }
    }

    fn toggle(&mut self) {
        if self.recording {
            eprintln!(
                "Recording stopped: {} frames saved to {}",
                self.frame_count,
                self.dir.display()
            );
            self.recording = false;
        } else {
            self.session += 1;
            self.dir = PathBuf::from(format!("recording_{:03}", self.session));
            fs::create_dir_all(&self.dir).expect("failed to create recording directory");
            self.frame_count = 0;
            self.recording = true;
            eprintln!("Recording started: {}", self.dir.display());
        }
    }

    fn capture_frame(&mut self) {
        if !self.recording {
            return;
        }
        let path = self.dir.join(format!("frame_{:06}.png", self.frame_count));
        save_screenshot_to(&path);
        self.frame_count += 1;
    }

    fn save_single(&self) {
        let path = PathBuf::from(format!("screenshot_{}.png", timestamp_tag()));
        save_screenshot_to(&path);
        eprintln!("Screenshot saved: {}", path.display());
    }
}

fn save_screenshot_to(path: &std::path::Path) {
    let img = get_screen_data();
    let w = img.width as u32;
    let h = img.height as u32;
    if let Some(buffer) = image::ImageBuffer::<image::Rgba<u8>, _>::from_raw(w, h, img.bytes) {
        // Macroquad screen data is top-to-bottom, but may need vertical flip
        let flipped = image::imageops::flip_vertical(&buffer);
        if let Err(e) = flipped.save(path) {
            eprintln!("Failed to save {}: {e}", path.display());
        }
    }
}

/// Simple timestamp tag for screenshot filenames.
fn timestamp_tag() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

// ── Rendering helpers ───────────────────────────────────────────────

fn draw_ground_grid(center: Vec3, half_size: f32, spacing: f32) {
    let x0 = (center.x / spacing).floor() * spacing - half_size;
    let y0 = (center.y / spacing).floor() * spacing - half_size;
    let steps = (half_size * 2.0 / spacing) as i32;
    let color = Color::new(0.3, 0.3, 0.3, 0.4);
    for i in 0..=steps {
        let ox = x0 + i as f32 * spacing;
        let oy = y0 + i as f32 * spacing;
        draw_line_3d(Vec3::new(ox, y0, 0.0), Vec3::new(ox, y0 + half_size * 2.0, 0.0), color);
        draw_line_3d(Vec3::new(x0, oy, 0.0), Vec3::new(x0 + half_size * 2.0, oy, 0.0), color);
    }
}

fn segment_activation(sim: &WormSim, idx: usize) -> f32 {
    let seg = &sim.body.segments[idx];
    (seg.dorsal_activation + seg.ventral_activation
        + seg.left_activation + seg.right_activation)
        / 4.0
}

fn heat_color(t: f32) -> Color {
    let t = t.clamp(0.0, 1.0);
    Color::new(t, 0.25, 1.0 - t, 1.0)
}

fn draw_worm_body(sim: &WormSim) {
    for i in 0..SEGMENT_COUNT {
        let pos = sim.body.segments[i].position;
        let center = Vec3::new(pos[0], pos[1], pos[2]);
        let act = segment_activation(sim, i);
        let color = heat_color(act);
        let radius = if i == 0 { 0.4 } else { 0.3 };
        draw_sphere(center, radius, None, color);

        if i > 0 {
            let prev = sim.body.segments[i - 1].position;
            let prev_center = Vec3::new(prev[0], prev[1], prev[2]);
            draw_line_3d(prev_center, center, Color::new(0.7, 0.7, 0.7, 0.6));
        }
    }

    // Head direction indicator
    let head = &sim.body.segments[0];
    let dir = head.direction();
    let head_pos = Vec3::new(head.position[0], head.position[1], head.position[2]);
    let tip = head_pos + Vec3::new(dir[0], dir[1], dir[2]) * 0.8;
    draw_line_3d(head_pos, tip, YELLOW);
}

fn draw_food(sim: &WormSim, frame: u64) {
    let food = sim.body.environment.food_source;
    let center = Vec3::new(food[0], food[1], food[2]);
    let pulse = 0.5 + 0.1 * (frame as f32 * 0.05).sin();
    draw_sphere(center, pulse, None, Color::new(0.1, 0.9, 0.2, 0.8));
}

fn draw_brain_neurons(sim: &WormSim) {
    let com = sim.body.center_of_mass();
    let brain_offset = Vec3::new(com[0], com[1], com[2] + 8.0);

    for neuron in &sim.brain.cascade.neurons {
        let pos = neuron.soma.position;
        let world = brain_offset + Vec3::new(pos[0], pos[1], pos[2]);

        let (base_color, base_r) = if neuron.nuclei.is_sensory() {
            (Color::new(0.2, 0.9, 0.2, 0.9), 0.12)
        } else if neuron.nuclei.is_motor() {
            (Color::new(0.9, 0.2, 0.2, 0.9), 0.12)
        } else {
            (Color::new(0.3, 0.5, 0.9, 0.6), 0.08)
        };

        let fired = neuron.trace > 0;
        let r = if fired { base_r * 1.8 } else { base_r };
        let color = if fired {
            Color::new(
                (base_color.r + 0.3).min(1.0),
                (base_color.g + 0.3).min(1.0),
                (base_color.b + 0.3).min(1.0),
                1.0,
            )
        } else {
            base_color
        };

        draw_sphere(world, r, None, color);
    }

    // Brain bounding box
    let min = brain_offset + Vec3::new(-9.0, -1.5, 0.0);
    let max = brain_offset + Vec3::new(0.0, 1.5, 1.5);
    let c = Color::new(0.5, 0.5, 0.5, 0.2);
    draw_line_3d(Vec3::new(min.x, min.y, min.z), Vec3::new(max.x, min.y, min.z), c);
    draw_line_3d(Vec3::new(max.x, min.y, min.z), Vec3::new(max.x, max.y, min.z), c);
    draw_line_3d(Vec3::new(max.x, max.y, min.z), Vec3::new(min.x, max.y, min.z), c);
    draw_line_3d(Vec3::new(min.x, max.y, min.z), Vec3::new(min.x, min.y, min.z), c);
    draw_line_3d(Vec3::new(min.x, min.y, max.z), Vec3::new(max.x, min.y, max.z), c);
    draw_line_3d(Vec3::new(max.x, min.y, max.z), Vec3::new(max.x, max.y, max.z), c);
    draw_line_3d(Vec3::new(max.x, max.y, max.z), Vec3::new(min.x, max.y, max.z), c);
    draw_line_3d(Vec3::new(min.x, max.y, max.z), Vec3::new(min.x, min.y, max.z), c);
    draw_line_3d(Vec3::new(min.x, min.y, min.z), Vec3::new(min.x, min.y, max.z), c);
    draw_line_3d(Vec3::new(max.x, min.y, min.z), Vec3::new(max.x, min.y, max.z), c);
    draw_line_3d(Vec3::new(max.x, max.y, min.z), Vec3::new(max.x, max.y, max.z), c);
    draw_line_3d(Vec3::new(min.x, max.y, min.z), Vec3::new(min.x, max.y, max.z), c);
}

fn draw_hud(sim: &WormSim, paused: bool, speed: u32, recording: bool) {
    let diag = sim.diagnostics();

    let phase_name = diag.phase.name();
    let phase_color = match diag.phase {
        DevelopmentalPhase::Genesis => YELLOW,
        DevelopmentalPhase::Exposure => GREEN,
        DevelopmentalPhase::Differentiation => ORANGE,
        DevelopmentalPhase::Crystallization => SKYBLUE,
    };

    let y = 20.0;
    let x = 10.0;
    let line_h = 18.0;
    let font_size = 16.0;

    draw_text(
        &format!("Phase: {phase_name}  frame {}/{}", diag.phase_frame, sim.config.phase.total_frames()),
        x, y, font_size, phase_color,
    );
    draw_text(
        &format!(
            "S:{} M:{} I:{}  diff:{}/302",
            diag.sensory_count, diag.motor_count, diag.inter_count, diag.differentiated_count,
        ),
        x, y + line_h, font_size, WHITE,
    );
    draw_text(
        &format!(
            "food_dist: {:.1}  energy: {:.3}  distress: {:.3}",
            diag.distance_to_food, diag.energy, diag.distress,
        ),
        x, y + line_h * 2.0, font_size, WHITE,
    );
    draw_text(
        &format!("spikes: {}  speed: {}x", diag.total_spikes, speed),
        x, y + line_h * 3.0, font_size, WHITE,
    );

    if paused {
        draw_text("PAUSED", x, y + line_h * 4.0, font_size, RED);
    }

    if recording {
        draw_text("REC", screen_width() - 50.0, 20.0, 18.0, RED);
        draw_circle(screen_width() - 60.0, 15.0, 5.0, RED);
    }

    let screen_h = screen_height();
    draw_text(
        "Space=pause  Up/Down=speed  R=reset  F=follow  LMB=pan  RMB=orbit  Scroll=zoom  V=record  P=screenshot",
        x, screen_h - 10.0, 14.0, Color::new(0.6, 0.6, 0.6, 0.8),
    );
}

// ── Main loop ───────────────────────────────────────────────────────

#[macroquad::main("Elegans Visualizer")]
async fn main() {
    let config = WormConfig::default();
    let mut sim = WormSim::new(config);
    let mut camera = OrbitCamera::new();
    let mut recorder = Recorder::new();
    let mut paused = false;
    let mut speed: u32 = 1;
    let mut frame: u64 = 0;

    loop {
        // ── Input ──
        if is_key_pressed(KeyCode::Space) {
            paused = !paused;
        }
        if is_key_pressed(KeyCode::Up) {
            speed = (speed * 2).min(64);
        }
        if is_key_pressed(KeyCode::Down) {
            speed = (speed / 2).max(1);
        }
        if is_key_pressed(KeyCode::R) {
            camera.reset();
        }
        if is_key_pressed(KeyCode::F) {
            camera.follow = !camera.follow;
        }
        if is_key_pressed(KeyCode::V) {
            recorder.toggle();
        }
        if is_key_pressed(KeyCode::P) {
            recorder.save_single();
        }

        // ── Simulation step(s) ──
        if !paused {
            for _ in 0..speed {
                sim.tick();
                frame += 1;
            }
        }

        // ── Camera update ──
        let com = sim.body.center_of_mass();
        let worm_center = Vec3::new(com[0], com[1], com[2]);
        camera.update(worm_center);

        // ── 3D Rendering ──
        clear_background(Color::new(0.05, 0.05, 0.1, 1.0));

        set_camera(&camera.camera3d());

        draw_ground_grid(worm_center, 30.0, 2.0);
        draw_food(&sim, frame);
        draw_worm_body(&sim);
        draw_brain_neurons(&sim);

        // ── 2D HUD ──
        set_default_camera();
        draw_hud(&sim, paused, speed, recorder.recording);

        // ── Recording (after all rendering) ──
        recorder.capture_frame();

        next_frame().await;
    }
}