elegans 1.0.0

//! WormSim — coupled body + brain simulation.
//!
//! Seeds 302 genomically diverse neurons in 3D space matching the worm
//! body layout: excitatory pyramidals, inhibitory interneurons, relays,
//! gates, oscillator seeds, and modulators. Developmental phases drive
//! disc-based differentiation into sensory/motor roles.

use neuropool::{
    SpatialNeuron, SpatialSynapse, SpatialRuntime, SpatialRuntimeConfig, Axon, Nuclei,
    WiringConfig, wire_by_proximity,
};

use crate::body::{Environment, WormBody};
use crate::coupling::CouplingState;
use crate::disc::{
    self, DifferentiationState, DiscProgram,
    accumulate_spatial_pressure, accumulate_activity_pressure, differentiate,
};
use crate::phase::{DevelopmentalPhase, PhaseConfig, PhaseController};

/// Configuration for the worm simulation.
#[derive(Clone, Debug)]
pub struct WormConfig {
    /// Total neuron count.
    pub neuron_count: usize,
    /// Developmental phase durations.
    pub phase: PhaseConfig,
    /// Runtime configuration overrides.
    pub runtime: SpatialRuntimeConfig,
    /// Frame interval in microseconds.
    pub frame_interval_us: u64,
    /// Wiring configuration.
    pub wiring: WiringConfig,
    /// How often to rebuild coupling caches (frames).
    pub cache_rebuild_interval: u64,
    /// How often to apply developmental disc pressure (frames).
    pub disc_interval: u64,
    /// Activity correlation window for disc pressure (μs).
    pub disc_activity_window_us: u64,
}

impl Default for WormConfig {
    fn default() -> Self {
        Self {
            neuron_count: 302,
            phase: PhaseConfig::default(),
            runtime: {
                let mut rt = SpatialRuntimeConfig::default();
                // Keep default structural_interval (100 frames = 18 steps in 1800).
                // Only accelerate pruning so it happens within the developmental window.
                rt.pruning_interval = 10;        // soft pruning every 1000 frames
                rt.hard_prune_interval = 15;     // hard pruning every 1500 frames
                rt
            },
            frame_interval_us: 10_000, // 10ms per frame
            wiring: WiringConfig {
                max_distance: 3.0,  // local circuits, not global reach
                max_fanout: 8,      // selective targeting
                max_fanin: 20,
                default_magnitude: 50, // needs ~4 coincident inputs to fire
            },
            cache_rebuild_interval: 50,
            disc_interval: 10,
            disc_activity_window_us: 50_000,
        }
    }
}

/// Diagnostic snapshot of the simulation state.
#[derive(Clone, Debug)]
pub struct SimDiagnostics {
    pub phase: DevelopmentalPhase,
    pub phase_frame: u64,
    pub total_frame: u64,
    pub sensory_count: usize,
    pub motor_count: usize,
    pub inter_count: usize,
    pub differentiated_count: usize,
    pub total_spikes: u64,
    pub distance_to_food: f32,
    pub body_center: [f32; 3],
    pub energy: f32,
    pub distress: f32,
}

/// The coupled worm simulation.
pub struct WormSim {
    pub body: WormBody,
    pub brain: SpatialRuntime,
    pub coupling: CouplingState,
    pub discs: Vec<DiscProgram>,
    pub diff_state: DifferentiationState,
    pub phase: PhaseController,
    pub config: WormConfig,
    /// Previous energy for valence-gated consolidation (Δenergy).
    prev_energy: f32,
    /// EMA-smoothed energy delta for valence gating.
    delta_ema: f32,
}

impl WormSim {
    /// Create a new worm simulation with default settings.
    pub fn new(config: WormConfig) -> Self {
        let environment = Environment::default();
        let body = WormBody::new(environment);

        // Genomic seeding: diverse population with E/I balance,
        // spatial organization, and varied nuclei properties.
        let neurons = scatter_neurons(config.neuron_count);

        // Wire by proximity + commissure
        let mut synapses = initial_wiring(&neurons, &config.wiring);
        wire_commissure(&neurons, &mut synapses);
        synapses.rebuild_index(neurons.len());

        // Create runtime
        let brain = SpatialRuntime::new(neurons, synapses, config.runtime.clone());

        // Create coupling
        let mut coupling = CouplingState::new();
        coupling.rebuild_caches(&brain.cascade.neurons);

        // Create disc programs
        let discs = disc::elegans_discs();
        let diff_state = DifferentiationState::new(config.neuron_count, discs.len());

        let phase = PhaseController::new(config.phase.clone());

        let prev_energy = body.metabolism.energy;

        Self {
            body,
            brain,
            coupling,
            discs,
            diff_state,
            phase,
            config,
            prev_energy,
            delta_ema: 0.0,
        }
    }

    /// Run one simulation tick: sense → inject → cascade → read → actuate → physics.
    pub fn tick(&mut self) {
        // 1. Sense the body
        let snapshot = self.body.sense();

        // 2. Inject sensory signals into brain
        self.coupling.inject_sensory(&snapshot, &mut self.brain);

        // 3. Run brain cascade
        self.brain.step(self.config.frame_interval_us);

        // 4. Read motor output
        let cmd = self.coupling.read_motor(&self.brain);

        // 5. Actuate body
        self.body.actuate(&cmd);

        // 6. Body physics
        self.body.physics_step();

        // 6b. Fiber tract adaptation (myelination, sensitization)
        self.coupling.adapt();

        // 6c. Metabolic tick (energy drain, feeding, distress update)
        self.body.metabolic_tick();

        // 6d. Valence-gated consolidation: EMA-smoothed Δenergy modulates mastery budget.
        //
        // Raw per-tick deltas (~0.0004) are too noisy and small for threshold detection.
        // EMA smoothing (α=0.1, ~10-tick window) accumulates consistent trends before
        // triggering valence. Threshold 0.00005 fires after ~8-10 consistent ticks.
        //
        // SUSPENDED during differentiation: circuit under reconstruction.
        {
            let energy = self.body.metabolism.energy;
            let delta = energy - self.prev_energy;
            self.prev_energy = energy;

            // EMA smoothing: α=0.1 gives ~10-frame window
            self.delta_ema = 0.9 * self.delta_ema + 0.1 * delta;

            let base = self.config.runtime.mastery_budget_per_cycle;
            let valence = if self.phase.current_phase == DevelopmentalPhase::Differentiation {
                1.0
            } else if self.delta_ema > 0.00005 {
                2.0f32
            } else if self.delta_ema < -0.00005 {
                0.5
            } else {
                1.0
            };
            let phase_rate = self.phase.plasticity_rate();
            self.brain.set_mastery_budget((base as f32 * valence * phase_rate) as u32);
        }

        // 7. Developmental disc pressure (if in accumulation phase)
        let total_frame = self.phase.total_frame;
        if self.phase.discs_accumulate()
            && total_frame % self.config.disc_interval == 0
        {
            // Spatial pressure
            accumulate_spatial_pressure(
                &self.brain.cascade.neurons,
                &self.discs,
                &mut self.diff_state,
            );

            // Activity-correlated pressure
            let active_sensory = self.coupling.active_sensory_channels(&snapshot);
            let active_motor = self.coupling.active_motor_channels(&self.brain);
            accumulate_activity_pressure(
                &self.brain.cascade.neurons,
                &self.discs,
                &mut self.diff_state,
                &active_sensory,
                &active_motor,
                self.brain.time_us(),
                self.config.disc_activity_window_us,
            );
        }

        // 8. Differentiation (if in differentiation phase)
        if self.phase.discs_differentiate()
            && total_frame % self.config.disc_interval == 0
        {
            let count = differentiate(
                &mut self.brain.cascade.neurons,
                &self.discs,
                &mut self.diff_state,
                5, // gradual: max 5 neurons per round — spread disruption
            );
            if count > 0 {
                // Rebuild coupling caches when neurons differentiate
                self.coupling.rebuild_caches(&self.brain.cascade.neurons);
            }
        }

        // 9. Periodic coupling cache rebuild
        if total_frame % self.config.cache_rebuild_interval == 0 {
            self.coupling.rebuild_caches(&self.brain.cascade.neurons);
        }

        // 10. Advance phase
        self.phase.advance();
    }

    /// Run development through all phases.
    pub fn develop(&mut self) {
        let total = self.config.phase.total_frames();
        for _ in 0..total {
            self.tick();
        }
    }

    /// Run N additional ticks after development (for testing behavior).
    pub fn run(&mut self, frames: u64) {
        for _ in 0..frames {
            self.tick();
        }
    }

    /// Get diagnostic snapshot.
    pub fn diagnostics(&self) -> SimDiagnostics {
        let (s, m, i) = disc::count_roles(&self.brain.cascade.neurons);
        SimDiagnostics {
            phase: self.phase.current_phase,
            phase_frame: self.phase.phase_frame,
            total_frame: self.phase.total_frame,
            sensory_count: s,
            motor_count: m,
            inter_count: i,
            differentiated_count: self.diff_state.differentiated_count(),
            total_spikes: self.brain.cascade.total_spikes(),
            distance_to_food: self.body.distance_to_food(),
            body_center: self.body.center_of_mass(),
            energy: self.body.metabolism.energy,
            distress: self.body.metabolism.distress,
        }
    }
}

/// Seed a diverse population of neurons in the worm brain volume.
///
/// Genomic seeding — not 302 identical pyramidals, but a biologically
/// grounded initial population with E/I balance, varied nuclei properties,
/// and spatial organization reflecting the C. elegans body plan.
///
/// ## Population Mix
///
/// | Type | Ratio | Role |
/// |------|-------|------|
/// | Excitatory pyramidal | ~40% | Computational substrate |
/// | Inhibitory interneuron | ~25% | E/I balance from genesis |
/// | Relay | ~15% | Fast signal conduits |
/// | Gate | ~10% | Coincidence detectors |
/// | Oscillator | ~5% | Locomotion CPG seeds |
/// | Modulator | ~5% | Plasticity control |
///
/// ## Spatial Layout
///
/// Head-biased density (nerve ring analog):
/// - Head (x = 0 to -2.5): ~40% of neurons, dense cluster
/// - Body (x = -2.5 to -7): ~45% of neurons, moderate density
/// - Tail (x = -7 to -9): ~15% of neurons, sparse
///
/// Brain volume: x = [0, -9], y = [-1.5, 1.5], z = [0, 1.5].
fn scatter_neurons(count: usize) -> Vec<SpatialNeuron> {
    let mut neurons = Vec::with_capacity(count);
    let mut seed = 12345u64;

    // Population boundaries (cumulative out of 100)
    const PYRAMIDAL_END: u64 = 40;
    const INHIBITORY_END: u64 = 65;
    const RELAY_END: u64 = 80;
    const GATE_END: u64 = 90;
    const OSCILLATOR_END: u64 = 95;
    // Remaining 5% = modulator

    // Spatial region boundaries (cumulative neuron count ratios)
    let head_count = count * 40 / 100;
    let body_count = count * 45 / 100;
    // tail_count = count - head_count - body_count

    for i in 0..count {
        // --- Spatial position based on region ---
        let (x, y_spread, z_spread) = if i < head_count {
            // Head/nerve ring: dense cluster near x = 0 to -2.5
            seed = xorshift(seed);
            let x = -((seed % 1000) as f32 / 1000.0 * 2.5);
            (x, 1.5, 1.5)
        } else if i < head_count + body_count {
            // Body segments: moderate spread along x = -2.5 to -7
            seed = xorshift(seed);
            let x = -2.5 - ((seed % 1000) as f32 / 1000.0 * 4.5);
            (x, 1.2, 1.2)
        } else {
            // Tail: sparse, x = -7 to -9
            seed = xorshift(seed);
            let x = -7.0 - ((seed % 1000) as f32 / 1000.0 * 2.0);
            (x, 0.8, 0.8)
        };

        seed = xorshift(seed);
        let y = ((seed % 1000) as f32 / 1000.0 - 0.5) * 2.0 * y_spread;
        seed = xorshift(seed);
        let z = (seed % 1000) as f32 / 1000.0 * z_spread;

        // --- Nuclei type selection ---
        seed = xorshift(seed);
        let type_roll = seed % 100;

        let base_nuclei = if type_roll < PYRAMIDAL_END {
            Nuclei::pyramidal()
        } else if type_roll < INHIBITORY_END {
            Nuclei::interneuron()
        } else if type_roll < RELAY_END {
            Nuclei::relay()
        } else if type_roll < GATE_END {
            Nuclei::gate()
        } else if type_roll < OSCILLATOR_END {
            // Oscillator with varied periods (50-200ms = locomotion range)
            seed = xorshift(seed);
            let period = 50_000 + (seed % 150_000) as u32;
            Nuclei::oscillator(period)
        } else {
            // Modulator (chemical IDs 0-3 for different neuromodulator types)
            seed = xorshift(seed);
            Nuclei::modulator((seed % 4) as u8)
        };

        // --- Property variation: ±15% jitter on key nuclei fields ---
        let mut nuclei = base_nuclei;
        seed = xorshift(seed);
        nuclei.soma_size = jitter_u8(nuclei.soma_size, 15, &mut seed);
        nuclei.leak = jitter_u8(nuclei.leak, 15, &mut seed);
        nuclei.metabolic_rate = jitter_u8(nuclei.metabolic_rate, 15, &mut seed);
        nuclei.axon_affinity = jitter_u8(nuclei.axon_affinity, 15, &mut seed);
        nuclei.myelin_affinity = jitter_u8(nuclei.myelin_affinity, 15, &mut seed);

        let mut neuron = SpatialNeuron::at([x, y, z], nuclei);

        // --- Axon direction: region-dependent bias ---
        // Head neurons: axons project laterally and posteriorly (into body)
        // Body neurons: axons project along the ventral cord
        // Tail neurons: axons project anteriorly (toward body)
        seed = xorshift(seed);
        let ax_dx = if i < head_count {
            // Head: bias posterior (negative x)
            -1.0 + ((seed % 1000) as f32 / 1000.0 - 0.5) * 3.0
        } else if i >= head_count + body_count {
            // Tail: bias anterior (positive x)
            1.0 + ((seed % 1000) as f32 / 1000.0 - 0.5) * 2.0
        } else {
            // Body: mixed direction along cord
            ((seed % 1000) as f32 / 1000.0 - 0.5) * 4.0
        };
        seed = xorshift(seed);
        let ax_dy = ((seed % 1000) as f32 / 1000.0 - 0.5) * 2.0;
        seed = xorshift(seed);
        let ax_dz = ((seed % 1000) as f32 / 1000.0 - 0.5) * 1.0;
        neuron.axon = Axon::toward([
            neuron.soma.position[0] + ax_dx,
            neuron.soma.position[1] + ax_dy,
            neuron.soma.position[2] + ax_dz,
        ]);

        neurons.push(neuron);
    }

    neurons
}

/// Apply ±percent jitter to a u8 value using deterministic PRNG.
fn jitter_u8(value: u8, percent: u8, seed: &mut u64) -> u8 {
    *seed = xorshift(*seed);
    let range = (value as u64 * percent as u64) / 100;
    if range == 0 {
        return value;
    }
    let delta = (*seed % (range * 2 + 1)) as i64 - range as i64;
    (value as i64 + delta).clamp(1, 255) as u8
}

/// Initial proximity-based wiring.
fn initial_wiring(neurons: &[SpatialNeuron], config: &WiringConfig) -> neuropool::spatial::SpatialSynapseStore {
    wire_by_proximity(neurons, config)
}

/// Wire a commissure near the head — cross-lateral synapses for left↔right signaling.
///
/// The C. elegans nerve ring is a dense commissure at the head where left and
/// right neural bundles cross-connect. Without this, proximity wiring traps
/// signals in local ipsilateral loops, preventing contralateral motor responses
/// needed for steering.
///
/// Finds neurons in the head region (x > -2.5), splits by lateral position
/// (y < -0.1 = left, y > 0.1 = right), and creates bidirectional excitatory
/// synapses between the closest cross-lateral pairs.
fn wire_commissure(
    neurons: &[SpatialNeuron],
    synapses: &mut neuropool::spatial::SpatialSynapseStore,
) {
    const MAGNITUDE: u8 = 40;    // moderate — doesn't need to fire alone, supports convergent input
    const DELAY_US: u32 = 1_000; // 1ms commissure crossing time
    const MAX_PAIRS: usize = 20; // sparse bridge — enough for cross-lateral signaling

    // Head neurons split by lateral position
    let mut left: Vec<(u32, [f32; 3])> = Vec::new();
    let mut right: Vec<(u32, [f32; 3])> = Vec::new();

    for (idx, n) in neurons.iter().enumerate() {
        let pos = n.soma.position;
        // Head region only; skip midline neurons (|y| < 0.1)
        if pos[0] > -2.5 && pos[1].abs() > 0.1 {
            if pos[1] < 0.0 {
                left.push((idx as u32, pos));
            } else {
                right.push((idx as u32, pos));
            }
        }
    }

    if left.is_empty() || right.is_empty() {
        return;
    }

    // For each left neuron, find nearest right neuron
    let mut pairs: Vec<(u32, u32, f32)> = Vec::with_capacity(left.len());
    for &(l_idx, l_pos) in &left {
        let mut best_r = right[0].0;
        let mut best_dsq = f32::MAX;
        for &(r_idx, r_pos) in &right {
            let dx = l_pos[0] - r_pos[0];
            let dy = l_pos[1] - r_pos[1];
            let dz = l_pos[2] - r_pos[2];
            let dsq = dx * dx + dy * dy + dz * dz;
            if dsq < best_dsq {
                best_dsq = dsq;
                best_r = r_idx;
            }
        }
        pairs.push((l_idx, best_r, best_dsq));
    }

    // Closest cross-lateral pairs are most plausible commissure bridges
    pairs.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
    pairs.truncate(MAX_PAIRS);

    // Bidirectional excitatory synapses
    for (l, r, _) in pairs {
        synapses.add(SpatialSynapse::excitatory(l, r, MAGNITUDE, DELAY_US));
        synapses.add(SpatialSynapse::excitatory(r, l, MAGNITUDE, DELAY_US));
    }
}

/// Fast deterministic PRNG.
fn xorshift(mut state: u64) -> u64 {
    state ^= state << 13;
    state ^= state >> 7;
    state ^= state << 17;
    state
}

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

    #[test]
    fn sim_creation() {
        let sim = WormSim::new(WormConfig::default());
        assert_eq!(sim.brain.cascade.neurons.len(), 302);
        assert_eq!(sim.body.segments.len(), SEGMENT_COUNT);
        assert_eq!(sim.phase.current_phase, DevelopmentalPhase::Genesis);
    }

    #[test]
    fn sim_single_tick() {
        let mut sim = WormSim::new(WormConfig::default());
        sim.tick();
        assert_eq!(sim.phase.total_frame, 1);
    }

    #[test]
    fn sim_runs_genesis() {
        let config = WormConfig {
            phase: PhaseConfig {
                genesis_frames: 50,
                exposure_frames: 10,
                differentiation_frames: 10,
                crystallization_frames: 10,
            },
            ..Default::default()
        };
        let mut sim = WormSim::new(config);

        // Run through genesis
        for _ in 0..50 {
            sim.tick();
        }

        assert_eq!(sim.phase.current_phase, DevelopmentalPhase::Exposure);
    }

    #[test]
    fn sim_diagnostics() {
        let mut sim = WormSim::new(WormConfig::default());
        sim.tick();

        let diag = sim.diagnostics();
        assert_eq!(diag.sensory_count + diag.motor_count + diag.inter_count, 302);
        assert_eq!(diag.inter_count, 302); // all undifferentiated at start
    }

    #[test]
    fn scatter_produces_correct_count() {
        let neurons = scatter_neurons(302);
        assert_eq!(neurons.len(), 302);

        // Genomic seeding should produce a diverse population
        let excitatory = neurons.iter().filter(|n| n.nuclei.is_excitatory()).count();
        let inhibitory = neurons.iter().filter(|n| n.nuclei.is_inhibitory()).count();

        // Must have both excitatory and inhibitory neurons
        assert!(excitatory > 0, "no excitatory neurons");
        assert!(inhibitory > 0, "no inhibitory neurons");

        // Inhibitory should be roughly 25% (±10%)
        let inhib_pct = inhibitory as f32 / 302.0;
        assert!(
            inhib_pct > 0.15 && inhib_pct < 0.35,
            "inhibitory ratio out of range: {inhib_pct:.2} ({inhibitory}/302)",
        );

        // None should have sensory/motor interface at genesis
        for n in &neurons {
            assert!(!n.nuclei.is_sensory(), "no sensory at genesis");
            assert!(!n.nuclei.is_motor(), "no motor at genesis");
        }

        eprintln!("Genomic mix: excitatory={excitatory}, inhibitory={inhibitory}");
    }

    #[test]
    fn scatter_neurons_in_brain_volume() {
        let neurons = scatter_neurons(302);
        for n in &neurons {
            let pos = n.soma.position;
            assert!(pos[0] <= 0.0 && pos[0] >= -9.0, "x out of range: {}", pos[0]);
            assert!(pos[1] >= -1.5 && pos[1] <= 1.5, "y out of range: {}", pos[1]);
            assert!(pos[2] >= 0.0 && pos[2] <= 1.5, "z out of range: {}", pos[2]);
        }
    }

    #[test]
    fn initial_wiring_creates_connections() {
        let neurons = scatter_neurons(302);
        let config = WiringConfig {
            max_distance: 4.0,
            max_fanout: 12,
            max_fanin: 25,
            default_magnitude: 80,
        };
        let store = initial_wiring(&neurons, &config);
        assert!(store.len() > 0, "should have some synapses");
    }
}