//! Worm body — 3D physics with segments, muscles, touch, and chemosensing.
//!
//! A chain of segments in 3D space. Each segment has dorsal/ventral and
//! left/right muscle pairs (4 muscles per segment), touch sensors on all
//! sides, and proprioceptive bend sensing. The head segment has
//! chemosensory channels for detecting gradient direction toward food.
//!
//! 3D allows the worm to wiggle through soil, crawl over surfaces, and
//! coil — not just slide on a plane.

use std::f32::consts::PI;

/// Number of body segments (head to tail).
pub const SEGMENT_COUNT: usize = 10;

/// Obstacle detection radius per segment.
const TOUCH_RADIUS: f32 = 0.8;

/// Segment spacing (rest length).
const SEGMENT_SPACING: f32 = 1.0;

/// Spring stiffness for segment coupling.
const SPRING_K: f32 = 0.5;

/// Maximum angular velocity from muscle activation (radians per tick).
const MAX_ANGULAR_VEL: f32 = 0.15;

/// Forward propulsion per tick when muscles activate.
const PROPULSION_FORCE: f32 = 0.08;

/// Damping factor for angular velocity.
const ANGULAR_DAMPING: f32 = 0.85;

/// Gravity effect (downward pull on z-axis per tick).
const GRAVITY: f32 = 0.02;

/// Ground plane z-coordinate (worm rests on this).
const GROUND_Z: f32 = 0.0;

/// A single body segment in 3D space.
#[derive(Clone, Debug)]
pub struct Segment {
    /// Position in 3D space.
    pub position: [f32; 3],
    /// Heading — yaw angle in the XY plane (radians).
    pub yaw: f32,
    /// Pitch angle — tilt above/below horizontal (radians).
    pub pitch: f32,
    /// Dorsal muscle activation (0.0–1.0) — bends upward/toward dorsum.
    pub dorsal_activation: f32,
    /// Ventral muscle activation (0.0–1.0) — bends downward/toward ventrum.
    pub ventral_activation: f32,
    /// Left muscle activation (0.0–1.0) — turns left.
    pub left_activation: f32,
    /// Right muscle activation (0.0–1.0) — turns right.
    pub right_activation: f32,
    /// Touch sensors: [dorsal, ventral, left, right].
    pub touch: [bool; 4],
    /// Angular velocities [yaw_rate, pitch_rate].
    angular_velocity: [f32; 2],
}

/// Touch sensor indices.
pub const TOUCH_DORSAL: usize = 0;
pub const TOUCH_VENTRAL: usize = 1;
pub const TOUCH_LEFT: usize = 2;
pub const TOUCH_RIGHT: usize = 3;

impl Segment {
    fn new(position: [f32; 3], yaw: f32) -> Self {
        Self {
            position,
            yaw,
            pitch: 0.0,
            dorsal_activation: 0.0,
            ventral_activation: 0.0,
            left_activation: 0.0,
            right_activation: 0.0,
            touch: [false; 4],
            angular_velocity: [0.0, 0.0],
        }
    }

    /// Bend angle in the dorsal-ventral axis (proprioceptive signal).
    pub fn bend_dv(&self) -> f32 {
        self.dorsal_activation - self.ventral_activation
    }

    /// Bend angle in the left-right axis (proprioceptive signal).
    pub fn bend_lr(&self) -> f32 {
        self.left_activation - self.right_activation
    }

    /// Forward direction unit vector in 3D.
    pub fn direction(&self) -> [f32; 3] {
        let cp = self.pitch.cos();
        [
            self.yaw.cos() * cp,
            self.yaw.sin() * cp,
            self.pitch.sin(),
        ]
    }

    /// Dorsal direction (up from the worm's perspective).
    pub fn dorsal_dir(&self) -> [f32; 3] {
        let sp = self.pitch.sin();
        let cp = self.pitch.cos();
        [
            -self.yaw.cos() * sp,
            -self.yaw.sin() * sp,
            cp,
        ]
    }

    /// Left direction (perpendicular to forward, in the XY-dominated plane).
    pub fn left_dir(&self) -> [f32; 3] {
        [-self.yaw.sin(), self.yaw.cos(), 0.0]
    }
}

/// A 3D obstacle (sphere).
#[derive(Clone, Debug)]
pub struct Obstacle {
    pub center: [f32; 3],
    pub radius: f32,
}

/// The 3D environment the worm inhabits.
#[derive(Clone, Debug)]
pub struct Environment {
    /// Chemical gradient source (food) in 3D.
    pub food_source: [f32; 3],
    /// Obstacle spheres.
    pub obstacles: Vec<Obstacle>,
    /// Whether to apply gravity (worm falls if not on ground).
    pub gravity_enabled: bool,
}

impl Default for Environment {
    fn default() -> Self {
        Self {
            food_source: [20.0, 0.0, 0.0],
            obstacles: Vec::new(),
            gravity_enabled: true,
        }
    }
}

/// Metabolic state — internal energy budget.
///
/// Energy drains from existing (basal) and moving (movement cost).
/// Energy feeds from food proximity. Distress signal is derived from
/// low energy and reaches the brain through the interoceptive tract.
#[derive(Clone, Debug)]
pub struct MetabolicState {
    /// Energy level (0.0 = dead, 1.0 = full).
    pub energy: f32,
    /// Basal metabolic drain per tick (living costs energy).
    pub basal_rate: f32,
    /// Movement cost multiplier (cost = rate * total_muscle_activation).
    pub movement_cost_rate: f32,
    /// Feeding rate multiplier (feed = rate * food_concentration_at_head).
    pub feeding_rate: f32,
    /// Distress signal (0.0 = fine, 1.0 = starving). Derived from energy.
    pub distress: f32,
}

impl Default for MetabolicState {
    fn default() -> Self {
        Self {
            energy: 0.5,
            basal_rate: 0.0002,
            movement_cost_rate: 0.000005,
            feeding_rate: 0.0008,
            distress: 0.5,
        }
    }
}

/// The worm body — a chain of segments in 3D.
#[derive(Clone, Debug)]
pub struct WormBody {
    pub segments: Vec<Segment>,
    pub environment: Environment,
    pub metabolism: MetabolicState,
}

/// Sensory output from the body (raw values for fiber tract transmission).
#[derive(Clone, Debug)]
pub struct SensorySnapshot {
    /// Chemosensory: [left_gradient, right_gradient, dorsal_gradient, concentration]
    pub chemosensory: [i32; 4],
    /// Metabolic distress (0 = fine, 1000 = starving).
    pub distress: i32,
    /// Touch per segment: [dorsal, ventral, left, right] as intensities.
    pub touch: [[i32; 4]; SEGMENT_COUNT],
    /// Proprioceptive bend per segment: [dorsal-ventral, left-right].
    pub proprioception: [[i32; 2]; SEGMENT_COUNT],
}

/// Motor commands from the brain (after fiber tract shaping).
#[derive(Clone, Debug)]
pub struct MotorCommand {
    /// Dorsal muscle activation per segment (0.0–1.0).
    pub dorsal: [f32; SEGMENT_COUNT],
    /// Ventral muscle activation per segment (0.0–1.0).
    pub ventral: [f32; SEGMENT_COUNT],
    /// Left muscle activation per segment (0.0–1.0).
    pub left: [f32; SEGMENT_COUNT],
    /// Right muscle activation per segment (0.0–1.0).
    pub right: [f32; SEGMENT_COUNT],
}

impl Default for MotorCommand {
    fn default() -> Self {
        Self {
            dorsal: [0.0; SEGMENT_COUNT],
            ventral: [0.0; SEGMENT_COUNT],
            left: [0.0; SEGMENT_COUNT],
            right: [0.0; SEGMENT_COUNT],
        }
    }
}

/// Number of sensory channels per interface type.
pub const CHEMO_CHANNELS: usize = 4;
pub const DISTRESS_CHANNELS: usize = 1;
pub const TOUCH_CHANNELS_PER_SEGMENT: usize = 4;
pub const PROPRIO_CHANNELS_PER_SEGMENT: usize = 2;
pub const MOTOR_CHANNELS_PER_SEGMENT: usize = 4;

/// Total sensory channels.
pub const TOTAL_SENSORY: usize = CHEMO_CHANNELS + DISTRESS_CHANNELS
    + SEGMENT_COUNT * TOUCH_CHANNELS_PER_SEGMENT
    + SEGMENT_COUNT * PROPRIO_CHANNELS_PER_SEGMENT;

/// Total motor channels.
pub const TOTAL_MOTOR: usize = SEGMENT_COUNT * MOTOR_CHANNELS_PER_SEGMENT;

impl WormBody {
    /// Create a worm body at the origin, pointing in the +x direction.
    pub fn new(environment: Environment) -> Self {
        let segments = (0..SEGMENT_COUNT)
            .map(|i| {
                let x = -(i as f32) * SEGMENT_SPACING;
                Segment::new([x, 0.0, 0.0], 0.0)
            })
            .collect();

        Self {
            segments,
            environment,
            metabolism: MetabolicState::default(),
        }
    }

    /// Create a worm at a specific starting position and heading.
    pub fn at(start: [f32; 3], heading: f32, environment: Environment) -> Self {
        let dx = -heading.cos() * SEGMENT_SPACING;
        let dy = -heading.sin() * SEGMENT_SPACING;
        let segments = (0..SEGMENT_COUNT)
            .map(|i| {
                Segment::new(
                    [
                        start[0] + i as f32 * dx,
                        start[1] + i as f32 * dy,
                        start[2],
                    ],
                    heading,
                )
            })
            .collect();

        Self {
            segments,
            environment,
            metabolism: MetabolicState::default(),
        }
    }

    /// Head segment position.
    pub fn head_position(&self) -> [f32; 3] {
        self.segments[0].position
    }

    /// Distance from head to food source.
    pub fn distance_to_food(&self) -> f32 {
        distance_3d(self.head_position(), self.environment.food_source)
    }

    /// Read sensory state from the body.
    pub fn sense(&self) -> SensorySnapshot {
        let head = &self.segments[0];
        let food = self.environment.food_source;

        // Chemosensory: gradient direction relative to head orientation
        let to_food = sub_3d(food, head.position);
        let dist = length_3d(to_food).max(0.001);
        let food_dir = scale_3d(to_food, 1.0 / dist);

        // Project food direction onto head's local axes
        let left = head.left_dir();
        let dorsal = head.dorsal_dir();
        let forward = head.direction();

        let left_component = dot_3d(food_dir, left);
        let dorsal_component = dot_3d(food_dir, dorsal);
        let forward_component = dot_3d(food_dir, forward);

        // Concentration: inverse distance (scaled for i32 range)
        let concentration = (1000.0 / (1.0 + dist * 0.1)) as i32;

        let chemosensory = [
            (left_component * 500.0) as i32,           // left gradient
            (-left_component * 500.0) as i32,          // right gradient
            (dorsal_component * 500.0) as i32,         // dorsal gradient
            (forward_component * concentration as f32 / 500.0 * 500.0) as i32, // approach strength
        ];

        // Touch per segment
        let mut touch = [[0i32; 4]; SEGMENT_COUNT];
        for (i, seg) in self.segments.iter().enumerate() {
            let dorsal_pt = add_3d(seg.position, scale_3d(seg.dorsal_dir(), 0.5));
            let ventral_pt = sub_3d(seg.position, scale_3d(seg.dorsal_dir(), 0.5));
            let left_pt = add_3d(seg.position, scale_3d(seg.left_dir(), 0.5));
            let right_pt = sub_3d(seg.position, scale_3d(seg.left_dir(), 0.5));

            for obs in &self.environment.obstacles {
                let threshold = TOUCH_RADIUS + obs.radius;
                let dd = distance_3d(dorsal_pt, obs.center);
                let dv = distance_3d(ventral_pt, obs.center);
                let dl = distance_3d(left_pt, obs.center);
                let dr = distance_3d(right_pt, obs.center);

                if dd < threshold {
                    touch[i][TOUCH_DORSAL] = touch[i][TOUCH_DORSAL].max(
                        ((threshold - dd) / threshold * 1000.0) as i32,
                    );
                }
                if dv < threshold {
                    touch[i][TOUCH_VENTRAL] = touch[i][TOUCH_VENTRAL].max(
                        ((threshold - dv) / threshold * 1000.0) as i32,
                    );
                }
                if dl < threshold {
                    touch[i][TOUCH_LEFT] = touch[i][TOUCH_LEFT].max(
                        ((threshold - dl) / threshold * 1000.0) as i32,
                    );
                }
                if dr < threshold {
                    touch[i][TOUCH_RIGHT] = touch[i][TOUCH_RIGHT].max(
                        ((threshold - dr) / threshold * 1000.0) as i32,
                    );
                }
            }

            // Ground contact as ventral touch
            if seg.position[2] <= GROUND_Z + 0.3 {
                let ground_touch = ((0.3 - (seg.position[2] - GROUND_Z).max(0.0)) / 0.3 * 500.0) as i32;
                touch[i][TOUCH_VENTRAL] = touch[i][TOUCH_VENTRAL].max(ground_touch);
            }
        }

        // Proprioception: bend angles scaled to i32
        let mut proprioception = [[0i32; 2]; SEGMENT_COUNT];
        for (i, seg) in self.segments.iter().enumerate() {
            proprioception[i][0] = (seg.bend_dv() * 500.0) as i32;
            proprioception[i][1] = (seg.bend_lr() * 500.0) as i32;
        }

        // Metabolic distress scaled to i32 range
        let distress = (self.metabolism.distress * 1000.0) as i32;

        SensorySnapshot {
            chemosensory,
            distress,
            touch,
            proprioception,
        }
    }

    /// Apply motor commands to the body.
    pub fn actuate(&mut self, cmd: &MotorCommand) {
        for (i, seg) in self.segments.iter_mut().enumerate() {
            seg.dorsal_activation = cmd.dorsal[i].clamp(0.0, 1.0);
            seg.ventral_activation = cmd.ventral[i].clamp(0.0, 1.0);
            seg.left_activation = cmd.left[i].clamp(0.0, 1.0);
            seg.right_activation = cmd.right[i].clamp(0.0, 1.0);
        }
    }

    /// Run one physics step: muscles → torque → coupling → gravity → movement.
    pub fn physics_step(&mut self) {
        // Phase 1: Muscle activation → angular velocity
        for seg in &mut self.segments {
            // Left-right difference → yaw torque
            let yaw_torque = (seg.left_activation - seg.right_activation) * MAX_ANGULAR_VEL;
            seg.angular_velocity[0] += yaw_torque;
            seg.angular_velocity[0] *= ANGULAR_DAMPING;
            seg.yaw += seg.angular_velocity[0];

            // Dorsal-ventral difference → pitch torque
            let pitch_torque = (seg.dorsal_activation - seg.ventral_activation) * MAX_ANGULAR_VEL;
            seg.angular_velocity[1] += pitch_torque;
            seg.angular_velocity[1] *= ANGULAR_DAMPING;
            seg.pitch += seg.angular_velocity[1];
            // Clamp pitch to prevent full flips
            seg.pitch = seg.pitch.clamp(-PI / 3.0, PI / 3.0);

            // Wrap yaw
            while seg.yaw > PI { seg.yaw -= 2.0 * PI; }
            while seg.yaw < -PI { seg.yaw += 2.0 * PI; }
        }

        // Phase 2: Forward propulsion
        for seg in &mut self.segments {
            let total = seg.dorsal_activation + seg.ventral_activation
                + seg.left_activation + seg.right_activation;
            if total > 0.01 {
                let dir = seg.direction();
                let thrust = total * PROPULSION_FORCE * 0.5; // scale by 0.5 since 4 muscles now
                seg.position[0] += dir[0] * thrust;
                seg.position[1] += dir[1] * thrust;
                seg.position[2] += dir[2] * thrust;
            }
        }

        // Phase 3: Gravity
        if self.environment.gravity_enabled {
            for seg in &mut self.segments {
                if seg.position[2] > GROUND_Z {
                    seg.position[2] -= GRAVITY;
                    if seg.position[2] < GROUND_Z {
                        seg.position[2] = GROUND_Z;
                    }
                }
            }
        }

        // Phase 4: Segment coupling — spring forces maintain chain connectivity
        for i in 1..self.segments.len() {
            let prev_pos = self.segments[i - 1].position;
            let curr_pos = self.segments[i].position;

            let delta = sub_3d(prev_pos, curr_pos);
            let dist = length_3d(delta).max(0.001);

            let stretch = dist - SEGMENT_SPACING;
            if stretch.abs() > 0.01 {
                let force = stretch * SPRING_K;
                let f = scale_3d(delta, force / dist);

                self.segments[i].position[0] += f[0];
                self.segments[i].position[1] += f[1];
                self.segments[i].position[2] += f[2];

                // Orient follower toward predecessor (yaw only)
                let target_yaw = delta[1].atan2(delta[0]);
                let yaw_diff = angle_wrap(target_yaw - self.segments[i].yaw);
                self.segments[i].yaw += yaw_diff * 0.3;

                // Pitch alignment
                let target_pitch = (delta[2] / dist).asin();
                let pitch_diff = target_pitch - self.segments[i].pitch;
                self.segments[i].pitch += pitch_diff * 0.2;
            }
        }

        // Phase 5: Ground clamping
        for seg in &mut self.segments {
            if seg.position[2] < GROUND_Z {
                seg.position[2] = GROUND_Z;
            }
        }

        // Phase 6: Touch detection
        for seg in &mut self.segments {
            seg.touch = [false; 4];
        }

        for i in 0..self.segments.len() {
            let seg = &self.segments[i];
            let dorsal_pt = add_3d(seg.position, scale_3d(seg.dorsal_dir(), 0.5));
            let ventral_pt = sub_3d(seg.position, scale_3d(seg.dorsal_dir(), 0.5));
            let left_pt = add_3d(seg.position, scale_3d(seg.left_dir(), 0.5));
            let right_pt = sub_3d(seg.position, scale_3d(seg.left_dir(), 0.5));

            for obs in &self.environment.obstacles {
                let threshold = TOUCH_RADIUS + obs.radius;
                if distance_3d(dorsal_pt, obs.center) < threshold {
                    self.segments[i].touch[TOUCH_DORSAL] = true;
                }
                if distance_3d(ventral_pt, obs.center) < threshold {
                    self.segments[i].touch[TOUCH_VENTRAL] = true;
                }
                if distance_3d(left_pt, obs.center) < threshold {
                    self.segments[i].touch[TOUCH_LEFT] = true;
                }
                if distance_3d(right_pt, obs.center) < threshold {
                    self.segments[i].touch[TOUCH_RIGHT] = true;
                }
            }

            // Ground as ventral touch
            if self.segments[i].position[2] <= GROUND_Z + 0.1 {
                self.segments[i].touch[TOUCH_VENTRAL] = true;
            }
        }
    }

    /// Total body length (sum of inter-segment distances).
    pub fn body_length(&self) -> f32 {
        let mut len = 0.0;
        for i in 1..self.segments.len() {
            len += distance_3d(self.segments[i - 1].position, self.segments[i].position);
        }
        len
    }

    /// Center of mass.
    pub fn center_of_mass(&self) -> [f32; 3] {
        let n = self.segments.len() as f32;
        let mut c = [0.0f32; 3];
        for seg in &self.segments {
            c[0] += seg.position[0];
            c[1] += seg.position[1];
            c[2] += seg.position[2];
        }
        [c[0] / n, c[1] / n, c[2] / n]
    }

    /// Sum of all muscle activations across all segments.
    pub fn total_muscle_activation(&self) -> f32 {
        self.segments.iter().map(|s| {
            s.dorsal_activation + s.ventral_activation
                + s.left_activation + s.right_activation
        }).sum()
    }

    /// Food concentration at the head (0.0–1.0).
    /// Inverse distance: 1.0 when on food, decays with distance.
    pub fn food_concentration_at_head(&self) -> f32 {
        let dist = self.distance_to_food();
        1.0 / (1.0 + dist * 0.1)
    }

    /// Update metabolic state: drain energy, feed from food, derive distress.
    ///
    /// Call once per tick after physics_step.
    pub fn metabolic_tick(&mut self) {
        let total_activation = self.total_muscle_activation();
        let food_conc = self.food_concentration_at_head();

        // Drain: existing costs energy, moving costs more
        self.metabolism.energy -= self.metabolism.basal_rate;
        self.metabolism.energy -= self.metabolism.movement_cost_rate * total_activation;

        // Feed: food proximity restores energy
        self.metabolism.energy += self.metabolism.feeding_rate * food_conc;

        // Clamp
        self.metabolism.energy = self.metabolism.energy.clamp(0.0, 1.0);

        // Derive distress: low energy = high distress
        self.metabolism.distress = (1.0 - self.metabolism.energy).max(0.0);
    }
}

// =========================================================================
// 3D Vector Helpers
// =========================================================================

fn distance_3d(a: [f32; 3], b: [f32; 3]) -> f32 {
    length_3d(sub_3d(a, b))
}

fn length_3d(v: [f32; 3]) -> f32 {
    (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt()
}

fn sub_3d(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
    [a[0] - b[0], a[1] - b[1], a[2] - b[2]]
}

fn add_3d(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
    [a[0] + b[0], a[1] + b[1], a[2] + b[2]]
}

fn scale_3d(v: [f32; 3], s: f32) -> [f32; 3] {
    [v[0] * s, v[1] * s, v[2] * s]
}

fn dot_3d(a: [f32; 3], b: [f32; 3]) -> f32 {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

fn angle_wrap(mut a: f32) -> f32 {
    while a > PI { a -= 2.0 * PI; }
    while a < -PI { a += 2.0 * PI; }
    a
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn worm_creation_3d() {
        let body = WormBody::new(Environment::default());
        assert_eq!(body.segments.len(), SEGMENT_COUNT);
        // Head at origin
        let head = body.head_position();
        assert!(head[0].abs() < 0.001);
        assert!(head[1].abs() < 0.001);
        assert!(head[2].abs() < 0.001);
    }

    #[test]
    fn worm_at_position_3d() {
        let body = WormBody::at([5.0, 5.0, 1.0], 0.0, Environment {
            gravity_enabled: false,
            ..Default::default()
        });
        let head = body.head_position();
        assert!((head[0] - 5.0).abs() < 0.001);
        assert!((head[1] - 5.0).abs() < 0.001);
        assert!((head[2] - 1.0).abs() < 0.001);
    }

    #[test]
    fn segments_are_spaced_3d() {
        let body = WormBody::new(Environment::default());
        for i in 1..body.segments.len() {
            let d = distance_3d(body.segments[i - 1].position, body.segments[i].position);
            assert!((d - SEGMENT_SPACING).abs() < 0.01, "spacing: {d}");
        }
    }

    #[test]
    fn sensory_channels_complete() {
        let body = WormBody::new(Environment::default());
        let snap = body.sense();
        assert_eq!(snap.chemosensory.len(), 4);
        assert_eq!(snap.touch.len(), SEGMENT_COUNT);
        assert_eq!(snap.proprioception.len(), SEGMENT_COUNT);
        // Each touch has 4 directions
        assert_eq!(snap.touch[0].len(), 4);
        // Each proprioception has 2 axes
        assert_eq!(snap.proprioception[0].len(), 2);
    }

    #[test]
    fn chemosensory_detects_food_ahead() {
        let body = WormBody::new(Environment {
            food_source: [10.0, 0.0, 0.0], // directly ahead
            ..Default::default()
        });
        let snap = body.sense();
        // Concentration should be positive
        assert!(snap.chemosensory[3] > 0, "approach: {}", snap.chemosensory[3]);
    }

    #[test]
    fn chemosensory_detects_food_left() {
        let body = WormBody::new(Environment {
            food_source: [0.0, 10.0, 0.0], // food to the left
            ..Default::default()
        });
        let snap = body.sense();
        assert!(snap.chemosensory[0] > 0, "left gradient: {}", snap.chemosensory[0]);
    }

    #[test]
    fn four_muscle_movement() {
        let mut body = WormBody::new(Environment {
            gravity_enabled: false,
            ..Default::default()
        });
        let initial = body.head_position();

        // Symmetric activation → forward
        let mut cmd = MotorCommand::default();
        for i in 0..SEGMENT_COUNT {
            cmd.dorsal[i] = 0.3;
            cmd.ventral[i] = 0.3;
            cmd.left[i] = 0.3;
            cmd.right[i] = 0.3;
        }

        for _ in 0..100 {
            body.actuate(&cmd);
            body.physics_step();
        }

        let moved = distance_3d(initial, body.head_position());
        assert!(moved > 1.0, "worm should move forward: {moved}");
    }

    #[test]
    fn dorsal_ventral_turns_pitch() {
        let mut body = WormBody::new(Environment {
            gravity_enabled: false,
            ..Default::default()
        });

        let mut cmd = MotorCommand::default();
        cmd.dorsal[0] = 1.0; // dorsal only → pitch up

        for _ in 0..50 {
            body.actuate(&cmd);
            body.physics_step();
        }

        assert!(body.segments[0].pitch.abs() > 0.05, "should have pitched");
    }

    #[test]
    fn left_right_turns_yaw() {
        let mut body = WormBody::new(Environment {
            gravity_enabled: false,
            ..Default::default()
        });

        // Asymmetric left/right → yaw torque
        let mut cmd = MotorCommand::default();
        cmd.left[0] = 1.0;
        cmd.right[0] = 0.0;

        // Only 5 steps — the worm turns fast enough that 50 steps produces
        // ~6 full rotations, and yaw wrapping lands near zero.
        for _ in 0..5 {
            body.actuate(&cmd);
            body.physics_step();
        }

        // After 5 steps, yaw ≈ 1.57 radians (well above threshold)
        assert!(body.segments[0].yaw.abs() > 0.1, "should have yawed: {}", body.segments[0].yaw);
    }

    #[test]
    fn gravity_pulls_down() {
        let mut body = WormBody::at([0.0, 0.0, 5.0], 0.0, Environment::default());

        for _ in 0..200 {
            body.physics_step();
        }

        // Should have fallen toward ground
        assert!(body.segments[0].position[2] < 3.0, "z: {}", body.segments[0].position[2]);
    }

    #[test]
    fn ground_clamping() {
        let mut body = WormBody::new(Environment::default());
        // Already on ground
        for _ in 0..50 {
            body.physics_step();
        }
        // Should not go below ground
        for seg in &body.segments {
            assert!(seg.position[2] >= GROUND_Z, "below ground: {}", seg.position[2]);
        }
    }

    #[test]
    fn touch_detects_obstacle_3d() {
        let body = WormBody::new(Environment {
            food_source: [20.0, 0.0, 0.0],
            obstacles: vec![Obstacle { center: [0.0, 0.5, 0.0], radius: 0.5 }],
            gravity_enabled: false,
        });

        let snap = body.sense();
        let any_touch = snap.touch.iter().any(|t| t.iter().any(|&v| v > 0));
        assert!(any_touch, "should detect nearby obstacle");
    }

    #[test]
    fn proprioception_reflects_4_muscles() {
        let mut body = WormBody::new(Environment::default());
        let mut cmd = MotorCommand::default();
        cmd.dorsal[3] = 0.8;
        cmd.ventral[3] = 0.2;
        cmd.left[5] = 0.7;
        cmd.right[5] = 0.1;
        body.actuate(&cmd);

        let snap = body.sense();
        assert!(snap.proprioception[3][0] > 0, "dv bend: {}", snap.proprioception[3][0]);
        assert!(snap.proprioception[5][1] > 0, "lr bend: {}", snap.proprioception[5][1]);
    }

    #[test]
    fn spring_coupling_3d() {
        let mut body = WormBody::new(Environment {
            gravity_enabled: false,
            ..Default::default()
        });

        let mut cmd = MotorCommand::default();
        cmd.dorsal[0] = 1.0;
        cmd.ventral[0] = 1.0;
        cmd.left[0] = 1.0;
        cmd.right[0] = 1.0;

        for _ in 0..200 {
            body.actuate(&cmd);
            body.physics_step();
        }

        for i in 1..body.segments.len() {
            let d = distance_3d(body.segments[i - 1].position, body.segments[i].position);
            assert!(d < SEGMENT_SPACING * 3.0, "segments {}-{} too far: {d}", i-1, i);
        }
    }

    #[test]
    fn distance_to_food() {
        let body = WormBody::new(Environment {
            food_source: [10.0, 0.0, 0.0],
            ..Default::default()
        });
        let d = body.distance_to_food();
        assert!((d - 10.0).abs() < 0.1);
    }

    #[test]
    fn channel_counts() {
        // 4 chemo + 1 distress + 10*4 touch + 10*2 proprio = 65 sensory
        assert_eq!(TOTAL_SENSORY, 65);
        // 10*4 motor = 40
        assert_eq!(TOTAL_MOTOR, 40);
    }

    #[test]
    fn metabolism_drains_without_food() {
        let mut body = WormBody::new(Environment {
            food_source: [1000.0, 0.0, 0.0], // very far food
            gravity_enabled: false,
            ..Default::default()
        });
        let initial_energy = body.metabolism.energy;

        // Activate muscles and tick metabolism
        let mut cmd = MotorCommand::default();
        for i in 0..SEGMENT_COUNT {
            cmd.dorsal[i] = 0.5;
            cmd.ventral[i] = 0.5;
        }

        for _ in 0..200 {
            body.actuate(&cmd);
            body.physics_step();
            body.metabolic_tick();
        }

        assert!(
            body.metabolism.energy < initial_energy,
            "energy should drain: {} vs {}",
            body.metabolism.energy, initial_energy,
        );
        assert!(
            body.metabolism.distress > 0.5,
            "distress should be elevated: {}",
            body.metabolism.distress,
        );
    }

    #[test]
    fn metabolism_feeds_near_food() {
        let mut body = WormBody::new(Environment {
            food_source: [0.0, 0.0, 0.0], // food right at head
            gravity_enabled: false,
            ..Default::default()
        });
        // Start at low energy
        body.metabolism.energy = 0.2;

        for _ in 0..200 {
            body.physics_step();
            body.metabolic_tick();
        }

        assert!(
            body.metabolism.energy > 0.2,
            "energy should increase near food: {}",
            body.metabolism.energy,
        );
        assert!(
            body.metabolism.distress < 0.8,
            "distress should decrease near food: {}",
            body.metabolism.distress,
        );
    }

    #[test]
    fn distress_in_sensory_snapshot() {
        let body = WormBody::new(Environment::default());
        let snap = body.sense();
        // Default energy = 0.5 → distress = 0.5 → i32 = 500
        assert!(snap.distress > 0, "distress should be nonzero at half energy: {}", snap.distress);
        assert!(snap.distress <= 1000, "distress should be <= 1000: {}", snap.distress);
    }
}
