elegans 1.0.0

//! Imaginal disc programs — latent developmental rules.
//!
//! Discs are not objects attached to neurons. They are developmental rules
//! that run during specific phases. Each rule checks conditions on a neuron
//! and, if met, shifts its nuclei properties toward a target profile.
//!
//! "Roles should not be properties.
//!  Roles should be phase-dependent behaviors that emerge when conditions are met."
//!
//! A neuron doesn't get assigned "sensory" — it *becomes* sensory when:
//! 1. It's close to a sensory anchor point (spatial proximity)
//! 2. Its firing correlates with external sensory input (activity correlation)
//! 3. Enough pressure accumulates (threshold crossing)
//! 4. The developmental phase permits differentiation

use neuropool::{Interface, Nuclei, Polarity, SpatialNeuron};

/// What a disc program can differentiate a neuron into.
#[derive(Clone, Debug)]
pub enum DiscRole {
    /// Sensory neuron — transduces external input on a channel.
    Sensory { channel: u16, modality: u8 },
    /// Motor neuron — outputs to actuator on a channel.
    Motor { channel: u16, modality: u8 },
    /// Command interneuron — hub with high connectivity.
    CommandInterneuron,
    /// Oscillator — autonomous rhythm generator.
    Oscillator { period_us: u32 },
}

/// A single imaginal disc program.
///
/// Each disc defines a spatial anchor, influence radius, target role,
/// and the threshold of accumulated pressure needed to trigger
/// differentiation.
#[derive(Clone, Debug)]
pub struct DiscProgram {
    /// What this disc produces when activated.
    pub target_role: DiscRole,
    /// Spatial anchor — where in 3D space this disc's influence is strongest.
    pub anchor: [f32; 3],
    /// Maximum distance from anchor for any influence.
    pub radius: f32,
    /// Pressure threshold to trigger differentiation.
    pub threshold: f32,
}

impl DiscProgram {
    pub fn sensory(channel: u16, modality: u8, anchor: [f32; 3], radius: f32) -> Self {
        Self {
            target_role: DiscRole::Sensory { channel, modality },
            anchor,
            radius,
            threshold: 35.0, // spatial pressure alone sufficient for close neurons
        }
    }

    pub fn motor(channel: u16, modality: u8, anchor: [f32; 3], radius: f32) -> Self {
        Self {
            target_role: DiscRole::Motor { channel, modality },
            anchor,
            radius,
            threshold: 20.0, // lower bar — motor has no activity correlation during exposure
        }
    }

    pub fn command_interneuron(anchor: [f32; 3], radius: f32) -> Self {
        Self {
            target_role: DiscRole::CommandInterneuron,
            anchor,
            radius,
            threshold: 35.0, // needs proximity and moderate activity
        }
    }

    pub fn oscillator(period_us: u32, anchor: [f32; 3], radius: f32) -> Self {
        Self {
            target_role: DiscRole::Oscillator { period_us },
            anchor,
            radius,
            threshold: 30.0, // lowered for reduced-activity regime
        }
    }

    /// Spatial influence on a neuron (0.0 if outside radius, 1.0 at anchor).
    pub fn spatial_influence(&self, neuron_pos: [f32; 3]) -> f32 {
        let dx = neuron_pos[0] - self.anchor[0];
        let dy = neuron_pos[1] - self.anchor[1];
        let dz = neuron_pos[2] - self.anchor[2];
        let dist = (dx * dx + dy * dy + dz * dz).sqrt();

        if dist >= self.radius {
            0.0
        } else {
            1.0 - dist / self.radius
        }
    }

    /// Target nuclei preset for this disc's role.
    ///
    /// Sensory neurons get high leak (200) so they decay quickly without
    /// sustained injection. This makes them transient responders — driven
    /// by their sensory input, not by network oscillation. Mirrors real
    /// C. elegans amphid neurons which are receptor-adapted cells.
    fn target_nuclei(&self) -> Nuclei {
        match &self.target_role {
            DiscRole::Sensory { channel, modality } => {
                let mut n = Nuclei::sensory(*channel, *modality);
                n.leak = 200; // fast decay — input-driven, not network-driven
                n
            }
            DiscRole::Motor { channel, modality } => Nuclei::motor(*channel, *modality),
            DiscRole::CommandInterneuron => Nuclei {
                soma_size: 220,
                axon_affinity: 200,
                myelin_affinity: 180,
                metabolic_rate: 120,
                leak: 90,
                refractory: 1500,
                oscillation_period: 0,
                interface: Interface::none(),
                polarity: Polarity::Positive,
            },
            DiscRole::Oscillator { period_us } => Nuclei::oscillator(*period_us),
        }
    }
}

/// Tracks differentiation pressure per neuron per disc.
///
/// Each neuron accumulates pressure from each disc independently.
/// When pressure crosses the disc's threshold, differentiation occurs.
pub struct DifferentiationState {
    /// Pressure matrix: [neuron_idx][disc_idx].
    pressure: Vec<Vec<f32>>,
    /// Whether each neuron has already differentiated (irreversible).
    pub differentiated: Vec<bool>,
    /// Which disc differentiated each neuron (if any).
    pub committed_disc: Vec<Option<usize>>,
}

impl DifferentiationState {
    pub fn new(neuron_count: usize, disc_count: usize) -> Self {
        Self {
            pressure: vec![vec![0.0; disc_count]; neuron_count],
            differentiated: vec![false; neuron_count],
            committed_disc: vec![None; neuron_count],
        }
    }

    /// Get pressure for a specific neuron and disc.
    pub fn pressure(&self, neuron_idx: usize, disc_idx: usize) -> f32 {
        self.pressure[neuron_idx][disc_idx]
    }

    /// Accumulate pressure on a neuron from a disc.
    pub fn accumulate(&mut self, neuron_idx: usize, disc_idx: usize, amount: f32) {
        if self.differentiated[neuron_idx] {
            return; // already committed
        }
        self.pressure[neuron_idx][disc_idx] += amount;
    }

    /// Find the disc with highest pressure for a neuron (if above threshold).
    pub fn winning_disc(&self, neuron_idx: usize, discs: &[DiscProgram]) -> Option<usize> {
        if self.differentiated[neuron_idx] {
            return None; // already done
        }

        let mut best_disc = None;
        let mut best_pressure = 0.0f32;

        for (disc_idx, disc) in discs.iter().enumerate() {
            let p = self.pressure[neuron_idx][disc_idx];
            if p >= disc.threshold && p > best_pressure {
                best_disc = Some(disc_idx);
                best_pressure = p;
            }
        }

        best_disc
    }

    /// Count how many neurons have differentiated.
    pub fn differentiated_count(&self) -> usize {
        self.differentiated.iter().filter(|&&d| d).count()
    }

    /// Count neurons committed to each disc.
    pub fn disc_populations(&self, disc_count: usize) -> Vec<usize> {
        let mut counts = vec![0; disc_count];
        for disc_idx in self.committed_disc.iter().flatten() {
            counts[*disc_idx] += 1;
        }
        counts
    }
}

/// Run one accumulation step: spatial proximity drives pressure.
///
/// Call this during the Exposure and Differentiation phases.
pub fn accumulate_spatial_pressure(
    neurons: &[SpatialNeuron],
    discs: &[DiscProgram],
    state: &mut DifferentiationState,
) {
    for (n_idx, neuron) in neurons.iter().enumerate() {
        if state.differentiated[n_idx] {
            continue;
        }
        // Only undifferentiated (internal) neurons can be influenced
        if neuron.nuclei.is_sensory() || neuron.nuclei.is_motor() {
            continue;
        }

        for (d_idx, disc) in discs.iter().enumerate() {
            let influence = disc.spatial_influence(neuron.soma.position);
            if influence > 0.0 {
                state.accumulate(n_idx, d_idx, influence * 1.5);
            }
        }
    }
}

/// Run one accumulation step: activity correlation drives pressure.
///
/// Neurons that fire when their disc's interface would expect activity
/// accumulate extra pressure. This is how "neurons that fire when
/// sensory input arrives" become sensory neurons.
pub fn accumulate_activity_pressure(
    neurons: &[SpatialNeuron],
    discs: &[DiscProgram],
    state: &mut DifferentiationState,
    active_sensory_channels: &[u16],
    active_motor_channels: &[u16],
    current_time_us: u64,
    window_us: u64,
) {
    for (n_idx, neuron) in neurons.iter().enumerate() {
        if state.differentiated[n_idx] {
            continue;
        }
        if neuron.nuclei.is_sensory() || neuron.nuclei.is_motor() {
            continue;
        }

        // Did this neuron fire recently?
        let fired = neuron.last_spike_us > current_time_us.saturating_sub(window_us)
            && neuron.last_spike_us > 0;
        if !fired {
            continue;
        }

        for (d_idx, disc) in discs.iter().enumerate() {
            let spatial = disc.spatial_influence(neuron.soma.position);
            if spatial <= 0.0 {
                continue;
            }

            let correlation_boost = match &disc.target_role {
                DiscRole::Sensory { channel, .. } => {
                    if active_sensory_channels.contains(channel) { 1.0 } else { 0.0 }
                }
                DiscRole::Motor { channel, .. } => {
                    if active_motor_channels.contains(channel) { 1.0 } else { 0.0 }
                }
                DiscRole::CommandInterneuron => {
                    // Neurons with many recent activations get hub pressure
                    if neuron.trace > 10 { 0.5 } else { 0.0 }
                }
                DiscRole::Oscillator { .. } => {
                    // Neurons that fire periodically get oscillator pressure
                    // (heuristic: high trace = consistent firing)
                    if neuron.trace > 15 { 0.3 } else { 0.0 }
                }
            };

            if correlation_boost > 0.0 {
                state.accumulate(n_idx, d_idx, spatial * correlation_boost);
            }
        }
    }
}

/// Population caps per role category.
///
/// C. elegans has ~80 sensory, ~113 motor, ~109 interneuron. We cap
/// sensory and motor to prevent either from consuming too many neurons.
/// Interneurons are what's left — they don't need a cap.
struct RoleCaps {
    max_sensory: usize,
    max_motor: usize,
    max_command: usize,
    max_oscillator: usize,
}

impl RoleCaps {
    /// Compute caps from total neuron count.
    ///
    /// Minimum of 1 per role so unit tests with small populations still work.
    fn for_population(n: usize) -> Self {
        Self {
            max_sensory: (n * 30 / 100).max(1),
            max_motor: (n * 40 / 100).max(1),
            max_command: (n * 5 / 100).max(1),
            max_oscillator: (n * 3 / 100).max(1),
        }
    }
}

/// Apply differentiation: neurons that crossed threshold become their target role.
///
/// This physically changes the neuron's nuclei — it's not a label, it's a
/// structural transformation. Call during the Differentiation phase.
///
/// Population caps prevent any single role from consuming too many neurons.
/// Interneurons are preserved as the computational bulk.
///
/// `max_per_round` caps how many neurons differentiate per call, spreading
/// the disruption across multiple rounds so the network can adapt
/// incrementally. Pass 0 for unlimited.
///
/// Returns the number of neurons that differentiated this step.
pub fn differentiate(
    neurons: &mut [SpatialNeuron],
    discs: &[DiscProgram],
    state: &mut DifferentiationState,
    max_per_round: usize,
) -> usize {
    let caps = RoleCaps::for_population(neurons.len());

    // Count existing role populations (from prior differentiation rounds)
    let mut sensory_count = 0usize;
    let mut motor_count = 0usize;
    let mut command_count = 0usize;
    let mut oscillator_count = 0usize;
    for n in neurons.iter() {
        if n.nuclei.is_sensory() { sensory_count += 1; }
        else if n.nuclei.is_motor() { motor_count += 1; }
        else if n.nuclei.is_oscillator() { oscillator_count += 1; }
    }
    for disc_idx in state.committed_disc.iter().flatten() {
        if matches!(discs[*disc_idx].target_role, DiscRole::CommandInterneuron) {
            command_count += 1;
        }
    }

    let mut count = 0;
    let limit = if max_per_round == 0 { usize::MAX } else { max_per_round };

    for n_idx in 0..neurons.len() {
        if count >= limit {
            break;
        }
        if let Some(disc_idx) = state.winning_disc(n_idx, discs) {
            let disc = &discs[disc_idx];

            // Check population cap for this role category
            let allowed = match &disc.target_role {
                DiscRole::Sensory { .. } => sensory_count < caps.max_sensory,
                DiscRole::Motor { .. } => motor_count < caps.max_motor,
                DiscRole::CommandInterneuron => command_count < caps.max_command,
                DiscRole::Oscillator { .. } => oscillator_count < caps.max_oscillator,
            };
            if !allowed {
                continue; // cap reached — preserve as interneuron
            }

            let target = disc.target_nuclei();

            // Gradual transition: blend current nuclei toward target
            let neuron = &mut neurons[n_idx];
            blend_nuclei(&mut neuron.nuclei, &target, 0.5);

            // Set interface (this is the "become" moment)
            neuron.nuclei.interface = target.interface;

            // Mark as differentiated
            state.differentiated[n_idx] = true;
            state.committed_disc[n_idx] = Some(disc_idx);
            count += 1;

            // Update running counts
            match &disc.target_role {
                DiscRole::Sensory { .. } => sensory_count += 1,
                DiscRole::Motor { .. } => motor_count += 1,
                DiscRole::CommandInterneuron => command_count += 1,
                DiscRole::Oscillator { .. } => oscillator_count += 1,
            }
        }
    }

    count
}

/// Blend nuclei properties from current toward target by factor (0.0–1.0).
///
/// Polarity is preserved for inhibitory neurons — GABAergic cells don't
/// become glutamatergic just because they differentiate near a sensory
/// interface. They can serve as inhibitory sensory/motor neurons.
fn blend_nuclei(current: &mut Nuclei, target: &Nuclei, factor: f32) {
    current.soma_size = lerp_u8(current.soma_size, target.soma_size, factor);
    current.axon_affinity = lerp_u8(current.axon_affinity, target.axon_affinity, factor);
    current.myelin_affinity = lerp_u8(current.myelin_affinity, target.myelin_affinity, factor);
    current.metabolic_rate = lerp_u8(current.metabolic_rate, target.metabolic_rate, factor);
    current.leak = lerp_u8(current.leak, target.leak, factor);
    current.refractory = lerp_u32(current.refractory, target.refractory, factor);

    // Preserve inhibitory polarity — neurotransmitter type is genomic,
    // not overridden by disc differentiation. Only modulators (Zero)
    // adopt the target's polarity.
    if current.polarity == Polarity::Zero {
        current.polarity = target.polarity;
    }

    if target.oscillation_period > 0 {
        current.oscillation_period = target.oscillation_period;
    }
}

fn lerp_u8(a: u8, b: u8, t: f32) -> u8 {
    let result = a as f32 * (1.0 - t) + b as f32 * t;
    result.round().clamp(0.0, 255.0) as u8
}

fn lerp_u32(a: u32, b: u32, t: f32) -> u32 {
    let result = a as f32 * (1.0 - t) + b as f32 * t;
    result.round().max(0.0) as u32
}

/// Count roles in the neuron population.
pub fn count_roles(neurons: &[SpatialNeuron]) -> (usize, usize, usize) {
    let mut sensory = 0;
    let mut motor = 0;
    let mut inter = 0;
    for n in neurons {
        if n.nuclei.is_sensory() {
            sensory += 1;
        } else if n.nuclei.is_motor() {
            motor += 1;
        } else {
            inter += 1;
        }
    }
    (sensory, motor, inter)
}

// =========================================================================
// Worm-specific disc programs
// =========================================================================

/// Modality constant for mechanoreceptive (touch) interfaces.
pub const MODALITY_MECHANO: u8 = 10;
/// Modality constant for chemosensory interfaces.
pub const MODALITY_CHEMO: u8 = 11;
/// Modality constant for proprioceptive interfaces.
pub const MODALITY_PROPRIO: u8 = 12;
/// Modality constant for motor interfaces.
pub const MODALITY_MOTOR: u8 = 13;

/// Create the imaginal disc programs for a C. elegans worm.
///
/// The worm brain volume is laid out along the x-axis (head at +x, tail at -x),
/// matching the body segment positions. Disc anchors are placed near body
/// interface points so that nearby neurons feel the pressure.
///
/// Layout:
/// - x = [0, 10]: head (+x) to tail (-x), matching segment positions
/// - y = [-2, 2]: lateral spread
/// - z = [0, 2]: dorsal-ventral
pub fn elegans_discs() -> Vec<DiscProgram> {
    let mut discs = Vec::new();
    let mut channel = 0u16;

    // Head chemosensory discs (4 channels near head position x~0)
    for i in 0..4u16 {
        let y_offset = (i as f32 - 1.5) * 0.5;
        discs.push(DiscProgram::sensory(
            channel, MODALITY_CHEMO,
            [0.5, y_offset, 1.0],
            1.5, // tighter radius — only neurons right at the head
        ));
        channel += 1;
    }

    // Touch sensory discs (4 per segment × 10 segments = 40 channels)
    for seg in 0..10u16 {
        let x = -(seg as f32); // segment positions go negative
        for dir in 0..4u16 {
            let (y_off, z_off) = match dir {
                0 => (0.0, 1.5),   // dorsal
                1 => (0.0, -0.5),  // ventral
                2 => (-1.0, 0.5),  // left
                _ => (1.0, 0.5),   // right
            };
            discs.push(DiscProgram::sensory(
                channel, MODALITY_MECHANO,
                [x, y_off, z_off],
                1.5, // population caps prevent over-differentiation
            ));
            channel += 1;
        }
    }

    // Proprioceptive sensory discs (2 per segment × 10 = 20 channels)
    for seg in 0..10u16 {
        let x = -(seg as f32);
        // DV proprioception
        discs.push(DiscProgram::sensory(
            channel, MODALITY_PROPRIO,
            [x, 0.0, 0.5],
            1.5, // population caps prevent over-differentiation
        ));
        channel += 1;
        // LR proprioception
        discs.push(DiscProgram::sensory(
            channel, MODALITY_PROPRIO,
            [x, 0.0, 0.5],
            1.5,
        ));
        channel += 1;
    }

    // Motor discs (4 per segment × 10 segments = 40 channels)
    //
    // Motor disc anchors are placed INWARD from sensory disc anchors.
    // In C. elegans, motor neurons live in the ventral nerve cord (core),
    // while sensory neurons are at the body wall (periphery).
    // Separating them spatially prevents sensory discs from outcompeting
    // motor discs via activity correlation (sensory input is active early
    // in development, but motor output requires motor neurons to exist).
    let mut motor_channel = 0u16;
    for seg in 0..10u16 {
        let x = -(seg as f32);
        for dir in 0..4u16 {
            let (y_off, z_off) = match dir {
                0 => (0.0, 1.0),   // dorsal motor (inward from sensory at 1.5)
                1 => (0.0, 0.2),   // ventral motor (inside brain volume, not -0.5)
                2 => (-0.5, 0.5),  // left motor (inward from sensory at -1.0)
                _ => (0.5, 0.5),   // right motor (inward from sensory at 1.0)
            };
            discs.push(DiscProgram::motor(
                motor_channel, MODALITY_MOTOR,
                [x, y_off, z_off],
                1.5,
            ));
            motor_channel += 1;
        }
    }

    // Command interneuron discs — a few hub positions along the nerve ring
    // (C. elegans has a nerve ring around the pharynx at the head)
    for i in 0..4 {
        let angle = i as f32 * std::f32::consts::PI * 0.5;
        discs.push(DiscProgram::command_interneuron(
            [0.0, angle.cos() * 0.8, angle.sin() * 0.8 + 0.5],
            2.5,
        ));
    }

    // Oscillator discs — ventral nerve cord pattern generators
    // One per body half (anterior/posterior)
    discs.push(DiscProgram::oscillator(
        5_000, // 5ms period — fast body wave
        [-2.0, 0.0, 0.0],
        3.0,
    ));
    discs.push(DiscProgram::oscillator(
        5_000,
        [-7.0, 0.0, 0.0],
        3.0,
    ));

    discs
}

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

    #[test]
    fn disc_spatial_influence() {
        let disc = DiscProgram::sensory(0, 1, [0.0, 0.0, 0.0], 5.0);

        // At anchor: full influence
        assert!((disc.spatial_influence([0.0, 0.0, 0.0]) - 1.0).abs() < 0.001);

        // At half radius: half influence
        assert!((disc.spatial_influence([2.5, 0.0, 0.0]) - 0.5).abs() < 0.001);

        // At radius: zero influence
        assert!((disc.spatial_influence([5.0, 0.0, 0.0])).abs() < 0.001);

        // Outside radius: zero
        assert_eq!(disc.spatial_influence([6.0, 0.0, 0.0]), 0.0);
    }

    #[test]
    fn accumulation_and_differentiation() {
        let discs = vec![
            DiscProgram::sensory(0, 1, [0.0, 0.0, 0.0], 5.0),
            DiscProgram::motor(0, 1, [10.0, 0.0, 0.0], 5.0),
        ];

        let mut neurons = vec![
            SpatialNeuron::pyramidal_at([0.5, 0.0, 0.0]), // near sensory disc
            SpatialNeuron::pyramidal_at([9.5, 0.0, 0.0]), // near motor disc
            SpatialNeuron::pyramidal_at([5.0, 0.0, 0.0]), // in between
        ];

        let mut state = DifferentiationState::new(3, 2);

        // Run many accumulation steps
        for _ in 0..200 {
            accumulate_spatial_pressure(&neurons, &discs, &mut state);
        }

        // Neuron 0 should have high pressure from disc 0
        assert!(state.pressure(0, 0) > state.pressure(0, 1));
        // Neuron 1 should have high pressure from disc 1
        assert!(state.pressure(1, 1) > state.pressure(1, 0));
        // Neuron 2 in between should have less pressure from both
        assert!(state.pressure(2, 0) < state.pressure(0, 0));

        // Differentiate
        let count = differentiate(&mut neurons, &discs, &mut state, 0);
        assert!(count >= 2, "at least 2 neurons should differentiate");

        // Check roles
        assert!(neurons[0].nuclei.is_sensory(), "near sensory disc should become sensory");
        assert!(neurons[1].nuclei.is_motor(), "near motor disc should become motor");
    }

    #[test]
    fn already_differentiated_neurons_skip() {
        let discs = vec![
            DiscProgram::sensory(0, 1, [0.0, 0.0, 0.0], 5.0),
        ];

        let neurons = vec![
            SpatialNeuron::sensory_at([0.0, 0.0, 0.0], 0, 1), // already sensory
        ];

        let mut state = DifferentiationState::new(1, 1);

        accumulate_spatial_pressure(&neurons, &discs, &mut state);

        // Existing sensory neuron should not accumulate pressure
        assert_eq!(state.pressure(0, 0), 0.0);
    }

    #[test]
    fn irreversible_differentiation() {
        let discs = vec![
            DiscProgram::sensory(0, 1, [0.0, 0.0, 0.0], 5.0),
            DiscProgram::motor(0, 1, [0.0, 0.0, 0.0], 5.0),
        ];

        let mut neurons = vec![SpatialNeuron::pyramidal_at([0.0, 0.0, 0.0])];
        let mut state = DifferentiationState::new(1, 2);

        // Pump pressure on disc 0 until differentiation
        for _ in 0..200 {
            accumulate_spatial_pressure(&neurons, &discs, &mut state);
        }
        differentiate(&mut neurons, &discs, &mut state, 0);

        assert!(state.differentiated[0]);

        // Further accumulation should not change anything
        let old_pressure = state.pressure(0, 1);
        accumulate_spatial_pressure(&neurons, &discs, &mut state);
        assert_eq!(state.pressure(0, 1), old_pressure);
    }

    #[test]
    fn elegans_disc_count() {
        let discs = elegans_discs();
        // 4 chemo + 40 touch + 20 proprio + 40 motor + 4 command + 2 oscillator = 110
        assert_eq!(discs.len(), 110);
    }

    #[test]
    fn count_roles_works() {
        let neurons = vec![
            SpatialNeuron::sensory_at([0.0, 0.0, 0.0], 0, 1),
            SpatialNeuron::motor_at([1.0, 0.0, 0.0], 0, 1),
            SpatialNeuron::pyramidal_at([2.0, 0.0, 0.0]),
            SpatialNeuron::pyramidal_at([3.0, 0.0, 0.0]),
        ];
        let (s, m, i) = count_roles(&neurons);
        assert_eq!(s, 1);
        assert_eq!(m, 1);
        assert_eq!(i, 2);
    }
}