elegans 1.0.0

//! POC validation tests for the elegans worm simulation.
//!
//! These tests validate that 302 undifferentiated neurons can develop
//! into a functional C. elegans nervous system through imaginal disc
//! developmental phases — with zero hardcoded structure.

use crate::body::Obstacle;
use crate::disc::count_roles;
use crate::phase::DevelopmentalPhase;
use crate::sim::{WormSim, WormConfig};
use crate::phase::PhaseConfig;

/// Helper: create a sim with shortened phases for testing.
fn test_sim() -> WormSim {
    WormSim::new(WormConfig {
        phase: PhaseConfig {
            genesis_frames: 200,
            exposure_frames: 800,
            differentiation_frames: 500,
            crystallization_frames: 300,
        },
        ..Default::default()
    })
}

// =========================================================================
// Structural Tests — Does development produce the right structure?
// =========================================================================

#[test]
fn development_completes() {
    let mut sim = test_sim();
    sim.develop();
    assert!(sim.phase.is_mature());
}

#[test]
fn neurons_differentiate_during_development() {
    let mut sim = test_sim();
    sim.develop();

    let diag = sim.diagnostics();
    assert!(
        diag.differentiated_count > 0,
        "no neurons differentiated: {}",
        diag.differentiated_count,
    );
}

#[test]
fn role_distribution_emerges() {
    let mut sim = test_sim();
    sim.develop();

    let (sensory, motor, inter) = count_roles(&sim.brain.cascade.neurons);

    // Some neurons should have become sensory
    assert!(sensory > 0, "no sensory neurons emerged");
    // Some neurons should have become motor
    assert!(motor > 0, "no motor neurons emerged");
    // Most should remain interneurons
    assert!(inter > 0, "no interneurons remain");
    // Total should still be 302
    assert_eq!(sensory + motor + inter, 302);

    eprintln!("Role distribution: sensory={sensory}, motor={motor}, inter={inter}");
}

#[test]
fn signal_flow_exists() {
    let mut sim = test_sim();
    sim.develop();

    // Run some ticks post-development and count spikes
    let spikes_before = sim.brain.cascade.total_spikes();
    sim.run(100);
    let spikes_after = sim.brain.cascade.total_spikes();

    assert!(
        spikes_after > spikes_before,
        "no neural activity post-development: {} → {}",
        spikes_before, spikes_after,
    );
}

#[test]
fn phase_transitions_occur() {
    let mut sim = test_sim();

    // Run through genesis
    let genesis_frames = sim.config.phase.genesis_frames;
    for _ in 0..genesis_frames {
        sim.tick();
    }
    assert_eq!(sim.phase.current_phase, DevelopmentalPhase::Exposure);

    // Run through exposure
    let exposure_frames = sim.config.phase.exposure_frames;
    for _ in 0..exposure_frames {
        sim.tick();
    }
    assert_eq!(sim.phase.current_phase, DevelopmentalPhase::Differentiation);
}

// =========================================================================
// Behavioral Tests — Does the worm behave like a worm?
// =========================================================================

#[test]
fn body_moves_during_development() {
    let mut sim = test_sim();
    let initial_pos = sim.body.center_of_mass();

    sim.develop();

    let final_pos = sim.body.center_of_mass();
    let moved = distance_3d(initial_pos, final_pos);

    // The worm should have moved at least a little (even random twitching)
    eprintln!("Body displacement during development: {moved:.3}");
    // Don't assert > 0 since early random neurons may not produce movement
}

#[test]
fn touch_withdrawal_after_development() {
    let mut sim = test_sim();
    sim.develop();

    // Place an obstacle near the head
    sim.body.environment.obstacles.push(Obstacle {
        center: [1.0, 0.0, 0.0],
        radius: 0.5,
    });

    // Record head position
    let pre_pos = sim.body.head_position();

    // Run 200 ticks with obstacle present
    sim.run(200);

    let post_pos = sim.body.head_position();
    let displacement = distance_3d(pre_pos, post_pos);

    eprintln!("Touch withdrawal displacement: {displacement:.3}");
    // Movement indicates the nervous system is processing touch input
}

#[test]
fn chemotaxis_after_development() {
    let mut sim = test_sim();
    sim.develop();

    // Place food ahead of the worm
    sim.body.environment.food_source = [10.0, 0.0, 0.0];
    let initial_distance = sim.body.distance_to_food();

    // Run 500 ticks
    sim.run(500);

    let final_distance = sim.body.distance_to_food();

    eprintln!(
        "Chemotaxis: distance {initial_distance:.1} → {final_distance:.1} (delta: {:.1})",
        initial_distance - final_distance,
    );
}

#[test]
fn locomotion_pattern_after_development() {
    let mut sim = test_sim();
    sim.develop();

    // Record segment angles over time
    let mut angle_history: Vec<Vec<f32>> = Vec::new();

    for _ in 0..200 {
        sim.tick();
        let angles: Vec<f32> = sim.body.segments.iter().map(|s| s.yaw).collect();
        angle_history.push(angles);
    }

    // Check for any angular variation (sign of muscle activity)
    let head_angles: Vec<f32> = angle_history.iter().map(|a| a[0]).collect();
    let variance = variance_f32(&head_angles);

    eprintln!("Head angle variance over 200 frames: {variance:.6}");
}

// =========================================================================
// Diagnostics Test — Print full development report
// =========================================================================

#[test]
fn development_report() {
    let mut sim = test_sim();

    eprintln!("\n=== ELEGANS DEVELOPMENT REPORT ===\n");

    // Snapshot before development
    let diag = sim.diagnostics();
    eprintln!("Pre-development:");
    eprintln!("  Neurons: 302 (all undifferentiated)");
    eprintln!("  Synapses: {}", sim.brain.cascade.synapses.len());
    eprintln!("  Food distance: {:.1}", diag.distance_to_food);

    // Run development with periodic snapshots
    let total = sim.config.phase.total_frames();
    let snapshot_interval = total / 10;

    for frame in 0..total {
        sim.tick();

        if frame > 0 && frame % snapshot_interval == 0 {
            let d = sim.diagnostics();
            eprintln!(
                "  Frame {}/{}: phase={:?}, differentiated={}, spikes={}, food_dist={:.1}, energy={:.3}, distress={:.3}",
                d.total_frame, total, d.phase, d.differentiated_count, d.total_spikes, d.distance_to_food, d.energy, d.distress,
            );
        }
    }

    // Final report
    let diag = sim.diagnostics();
    let (s, m, i) = count_roles(&sim.brain.cascade.neurons);

    eprintln!("\nPost-development:");
    eprintln!("  Phase: {:?}", diag.phase);
    eprintln!("  Differentiated: {}/302", diag.differentiated_count);
    eprintln!("  Sensory: {s}");
    eprintln!("  Motor: {m}");
    eprintln!("  Interneuron: {i}");
    eprintln!("  Total spikes: {}", diag.total_spikes);
    eprintln!("  Food distance: {:.1}", diag.distance_to_food);
    eprintln!("  Energy: {:.3}, Distress: {:.3}", diag.energy, diag.distress);
    eprintln!("  Learning: str={} weak={} flip={} dorm={} awake={}",
        sim.brain.learning.strengthened,
        sim.brain.learning.weakened,
        sim.brain.learning.flipped,
        sim.brain.learning.dormant,
        sim.brain.learning.awakened,
    );
    eprintln!("  Structural: migrate={} tissue={} prune={} hard_prune={}",
        sim.brain.structural.migration_steps,
        sim.brain.structural.tissue_updates,
        sim.brain.structural.pruning_cycles,
        sim.brain.structural.hard_pruned,
    );

    // Disc population breakdown
    let pops = sim.diff_state.disc_populations(sim.discs.len());
    let assigned: usize = pops.iter().sum();
    eprintln!("  Disc assignments: {assigned} neurons committed to {} discs", sim.discs.len());

    eprintln!("\n=== END REPORT ===\n");
}

// =========================================================================
// Diagnostic Probes — Where is the motor path dying?
// =========================================================================

/// Efferent probe: Are motor tracts seeing any meaningful signal?
///
/// After development, runs 200 ticks and logs:
/// - How many motor neurons have trace > 0 per frame
/// - Total motor signal magnitude reaching the motor command
/// - Per-segment muscle activation values
#[test]
fn probe_efferent_motor_signal() {
    let mut sim = test_sim();
    sim.develop();

    let mut total_motor_magnitude: u64 = 0;
    let mut frames_with_motor_signal = 0u64;
    let mut max_activation: f32 = 0.0;
    let mut total_trace_positive = 0u64;
    let mut peak_trace: i8 = 0;

    for frame in 0..200 {
        // Run one tick
        let snapshot = sim.body.sense();
        sim.coupling.inject_sensory(&snapshot, &mut sim.brain);
        sim.brain.step(sim.config.frame_interval_us);

        // Probe: check motor neuron traces BEFORE reading motor
        let mut trace_positive_this_frame = 0u32;
        for seg_coupled in &sim.coupling.segments {
            for &nidx in &seg_coupled.motor_neuron_indices {
                let neuron = &sim.brain.cascade.neurons[nidx as usize];
                if neuron.trace > 0 {
                    trace_positive_this_frame += 1;
                    peak_trace = peak_trace.max(neuron.trace);
                }
            }
        }
        total_trace_positive += trace_positive_this_frame as u64;

        // Read motor output through coupling
        let cmd = sim.coupling.read_motor(&sim.brain);

        // Measure motor command magnitude
        let mut frame_magnitude = 0.0f32;
        for seg_idx in 0..crate::body::SEGMENT_COUNT {
            let seg_sum = cmd.dorsal[seg_idx] + cmd.ventral[seg_idx]
                + cmd.left[seg_idx] + cmd.right[seg_idx];
            frame_magnitude += seg_sum;
            max_activation = max_activation.max(cmd.dorsal[seg_idx]);
            max_activation = max_activation.max(cmd.ventral[seg_idx]);
            max_activation = max_activation.max(cmd.left[seg_idx]);
            max_activation = max_activation.max(cmd.right[seg_idx]);
        }

        if frame_magnitude > 0.001 {
            frames_with_motor_signal += 1;
            total_motor_magnitude += (frame_magnitude * 1000.0) as u64;
        }

        // Per-segment detail for first few frames with signal
        if frame < 5 || (frame_magnitude > 0.001 && frames_with_motor_signal <= 3) {
            eprintln!(
                "  Frame {frame}: trace+={trace_positive_this_frame}, motor_mag={frame_magnitude:.4}, \
                 d0={:.3} v0={:.3} l0={:.3} r0={:.3}",
                cmd.dorsal[0], cmd.ventral[0], cmd.left[0], cmd.right[0],
            );
        }

        sim.body.actuate(&cmd);
        sim.body.physics_step();
        sim.phase.advance();
    }

    eprintln!("\n=== EFFERENT PROBE RESULTS ===");
    eprintln!("  Frames with motor signal: {frames_with_motor_signal}/200");
    eprintln!("  Total motor magnitude (x1000): {total_motor_magnitude}");
    eprintln!("  Peak single-channel activation: {max_activation:.4}");
    eprintln!("  Total motor-neuron trace-positive samples: {total_trace_positive}");
    eprintln!("  Peak trace value: {peak_trace}");
    eprintln!("=== END EFFERENT PROBE ===\n");
}

/// Physics probe: Does a hardcoded traveling wave move the body?
///
/// Bypasses the brain entirely. Injects a sinusoidal dorsal-ventral
/// traveling wave (like a real worm's locomotion CPG output) and
/// measures whether the body physics produces forward displacement.
#[test]
fn probe_physics_traveling_wave() {
    use crate::body::{WormBody, MotorCommand, Environment, SEGMENT_COUNT};

    let mut body = WormBody::new(Environment {
        food_source: [100.0, 0.0, 0.0],
        obstacles: Vec::new(),
        gravity_enabled: true,
    });

    let initial_pos = body.center_of_mass();
    let initial_head = body.head_position();

    // Traveling wave parameters:
    // - Each segment is phase-shifted by 2π/SEGMENT_COUNT
    // - Dorsal and ventral are anti-phase (when dorsal contracts, ventral relaxes)
    // - Wave travels from head to tail → forward propulsion
    let wave_freq = 0.1; // cycles per frame (10 Hz at 100fps)
    let phase_offset = std::f32::consts::TAU / SEGMENT_COUNT as f32;

    let mut max_displacement: f32 = 0.0;

    for frame in 0..200 {
        let mut cmd = MotorCommand::default();
        let t = frame as f32 * wave_freq * std::f32::consts::TAU;

        for seg in 0..SEGMENT_COUNT {
            let phase = t - seg as f32 * phase_offset;
            let wave = phase.sin(); // -1.0 to 1.0

            // Dorsal-ventral anti-phase: drives sinusoidal body undulation
            if wave > 0.0 {
                cmd.dorsal[seg] = wave * 0.8;
                cmd.ventral[seg] = 0.0;
            } else {
                cmd.dorsal[seg] = 0.0;
                cmd.ventral[seg] = (-wave) * 0.8;
            }
        }

        body.actuate(&cmd);
        body.physics_step();

        let disp = distance_3d(initial_pos, body.center_of_mass());
        max_displacement = max_displacement.max(disp);

        if frame % 50 == 0 {
            let pos = body.center_of_mass();
            eprintln!(
                "  Frame {frame}: center=[{:.2}, {:.2}, {:.2}], disp={disp:.3}",
                pos[0], pos[1], pos[2],
            );
        }
    }

    let final_pos = body.center_of_mass();
    let final_disp = distance_3d(initial_pos, final_pos);
    let head_disp = distance_3d(initial_head, body.head_position());

    eprintln!("\n=== PHYSICS PROBE RESULTS ===");
    eprintln!("  Initial center: [{:.2}, {:.2}, {:.2}]", initial_pos[0], initial_pos[1], initial_pos[2]);
    eprintln!("  Final center:   [{:.2}, {:.2}, {:.2}]", final_pos[0], final_pos[1], final_pos[2]);
    eprintln!("  CoM displacement: {final_disp:.3}");
    eprintln!("  Head displacement: {head_disp:.3}");
    eprintln!("  Max displacement during wave: {max_displacement:.3}");
    eprintln!("=== END PHYSICS PROBE ===\n");

    // The traveling wave SHOULD produce forward displacement
    // If it doesn't, the body physics itself is broken
    assert!(
        max_displacement > 0.1,
        "traveling wave should move the body: max_disp={max_displacement:.4}",
    );
}

// =========================================================================
// Causal pathway probe — does chemo input causally influence motor output?
// =========================================================================

/// Probe: Does injecting a left-right asymmetric chemo signal produce
/// asymmetric motor output? This tests the actual sensory→motor causal chain.
///
/// Protocol:
/// 1. Develop the worm
/// 2. Run 2000 post-development ticks (extended re-learning period)
/// 3. Inject left-biased chemo for 100 ticks, measure left-right motor asymmetry
/// 4. Inject right-biased chemo for 100 ticks, measure left-right motor asymmetry
/// 5. Report whether the motor response direction follows the sensory bias
#[test]
fn probe_causal_pathway() {
    use crate::body::SEGMENT_COUNT;

    let mut sim = test_sim();
    sim.develop();

    // Extended post-development learning period
    eprintln!("\n=== CAUSAL PATHWAY PROBE ===\n");
    eprintln!("Running 2000 post-dev ticks for re-learning...");
    sim.run(2000);

    let food_dist = sim.body.distance_to_food();
    eprintln!("  After 2000 post-dev ticks: food_dist={food_dist:.1}");

    // === Diagnostic: trace the signal chain ===
    eprintln!("\n  --- Signal chain diagnostic ---");

    // Head sensory neuron indices (chemo channels 0-3)
    let head_sens = sim.coupling.head.sensory_neuron_indices.clone();
    eprintln!("  Head sensory neuron indices: {:?}", &head_sens[..head_sens.len().min(5)]);

    // Check what chemo values the body produces for left food
    sim.body.environment.food_source = [5.0, -10.0, 0.0];
    let snap_left = sim.body.sense();
    eprintln!("  Chemo values (food LEFT):  {:?}", snap_left.chemosensory);

    sim.body.environment.food_source = [5.0, 10.0, 0.0];
    let snap_right = sim.body.sense();
    eprintln!("  Chemo values (food RIGHT): {:?}", snap_right.chemosensory);

    // Check membrane potentials of chemo neurons BEFORE injection
    for ch in 0..head_sens.len().min(4) {
        let idx = head_sens[ch] as usize;
        let n = &sim.brain.cascade.neurons[idx];
        eprintln!("  Chemo neuron {} (idx {}): membrane={}, trace={}, is_sensory={}, pos={:?}",
            ch, idx, n.membrane, n.trace,
            n.nuclei.is_sensory(), n.soma.position);
    }

    // Run ONE tick with left food and check membrane change
    sim.body.environment.food_source = [5.0, -10.0, 0.0];
    sim.tick();
    eprintln!("\n  After 1 tick with food LEFT:");
    for ch in 0..head_sens.len().min(4) {
        let idx = head_sens[ch] as usize;
        let n = &sim.brain.cascade.neurons[idx];
        eprintln!("    Chemo neuron {} (idx {}): membrane={}, trace={}",
            ch, idx, n.membrane, n.trace);
    }

    // Count total spikes, sensory spikes, motor spikes
    let total_spikes = sim.brain.cascade.neurons.iter().filter(|n| n.trace > 0).count();
    let sensory_spikes = sim.brain.cascade.neurons.iter()
        .filter(|n| n.trace > 0 && n.nuclei.is_sensory()).count();
    let motor_spikes = sim.brain.cascade.neurons.iter()
        .filter(|n| n.trace > 0 && n.nuclei.is_motor()).count();
    eprintln!("    Spikes this tick: total={}, sensory={}, motor={}", total_spikes, sensory_spikes, motor_spikes);

    // Check segment motor neuron indices for first 2 segments
    for seg in 0..2 {
        let motor_idx = &sim.coupling.segments[seg].motor_neuron_indices;
        eprintln!("    Segment {} motor neurons: {:?}", seg, &motor_idx[..motor_idx.len().min(4)]);
        for &mi in motor_idx.iter().take(4) {
            let n = &sim.brain.cascade.neurons[mi as usize];
            eprintln!("      Motor neuron {}: membrane={}, trace={}, is_motor={}",
                mi, n.membrane, n.trace, n.nuclei.is_motor());
        }
    }

    // === Place food at extreme positions relative to worm for clear gradient ===
    // Use worm's current center to position food nearby on opposite sides
    let center = sim.body.center_of_mass();
    let near_left = [center[0] + 3.0, center[1] - 8.0, 0.0];
    let near_right = [center[0] + 3.0, center[1] + 8.0, 0.0];
    eprintln!("  Worm center: {:?}", center);
    eprintln!("  Food LEFT: {:?}, Food RIGHT: {:?}", near_left, near_right);

    eprintln!("\n  --- Phase A: food LEFT (500 ticks) ---");
    sim.body.environment.food_source = near_left;
    let mut left_motor_sum = 0.0f32;
    let mut right_motor_sum = 0.0f32;
    let mut chemo_fire_count = 0u32;

    for t in 0..500 {
        sim.tick();
        let cmd = sim.coupling.read_motor(&sim.brain);
        for seg in 0..SEGMENT_COUNT {
            left_motor_sum += cmd.left[seg];
            right_motor_sum += cmd.right[seg];
        }
        // Count chemo neuron fires
        for ch in 0..head_sens.len().min(4) {
            let idx = head_sens[ch] as usize;
            if sim.brain.cascade.neurons[idx].trace > 0 {
                chemo_fire_count += 1;
            }
        }
        if t == 0 || t == 99 || t == 499 {
            let snap = sim.body.sense();
            eprintln!("    tick {}: L_motor={:.2} R_motor={:.2} chemo={:?}", t,
                cmd.left.iter().sum::<f32>(), cmd.right.iter().sum::<f32>(),
                snap.chemosensory);
        }
    }
    eprintln!("    Chemo neuron fires over 500 ticks: {}", chemo_fire_count);
    eprintln!("    Left motor total:  {left_motor_sum:.2}");
    eprintln!("    Right motor total: {right_motor_sum:.2}");
    eprintln!("    Asymmetry (L-R):   {:.2}", left_motor_sum - right_motor_sum);

    eprintln!("\n  --- Phase B: food RIGHT (500 ticks) ---");
    sim.body.environment.food_source = near_right;
    let mut left_motor_sum_b = 0.0f32;
    let mut right_motor_sum_b = 0.0f32;
    let mut chemo_fire_count_b = 0u32;

    for t in 0..500 {
        sim.tick();
        let cmd = sim.coupling.read_motor(&sim.brain);
        for seg in 0..SEGMENT_COUNT {
            left_motor_sum_b += cmd.left[seg];
            right_motor_sum_b += cmd.right[seg];
        }
        for ch in 0..head_sens.len().min(4) {
            let idx = head_sens[ch] as usize;
            if sim.brain.cascade.neurons[idx].trace > 0 {
                chemo_fire_count_b += 1;
            }
        }
        if t == 0 || t == 99 || t == 499 {
            let snap = sim.body.sense();
            eprintln!("    tick {}: L_motor={:.2} R_motor={:.2} chemo={:?}", t,
                cmd.left.iter().sum::<f32>(), cmd.right.iter().sum::<f32>(),
                snap.chemosensory);
        }
    }
    eprintln!("    Chemo neuron fires over 500 ticks: {}", chemo_fire_count_b);
    eprintln!("    Left motor total:  {left_motor_sum_b:.2}");
    eprintln!("    Right motor total: {right_motor_sum_b:.2}");
    eprintln!("    Asymmetry (L-R):   {:.2}", left_motor_sum_b - right_motor_sum_b);

    // Check if the asymmetry FLIPPED between phases
    let asym_a = left_motor_sum - right_motor_sum;
    let asym_b = left_motor_sum_b - right_motor_sum_b;
    let flipped = (asym_a > 0.0 && asym_b < 0.0) || (asym_a < 0.0 && asym_b > 0.0);

    eprintln!("\n  Asymmetry flip: {}", if flipped { "YES — causal pathway exists!" } else { "NO — no sensory→motor causation yet" });
    eprintln!("=== END CAUSAL PATHWAY PROBE ===\n");
}

// =========================================================================
// Helpers
// =========================================================================

fn distance_3d(a: [f32; 3], b: [f32; 3]) -> f32 {
    let dx = a[0] - b[0];
    let dy = a[1] - b[1];
    let dz = a[2] - b[2];
    (dx * dx + dy * dy + dz * dz).sqrt()
}

fn variance_f32(values: &[f32]) -> f32 {
    if values.is_empty() { return 0.0; }
    let mean = values.iter().sum::<f32>() / values.len() as f32;
    values.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / values.len() as f32
}