elegans 1.0.0

//! Brain↔body coupling via fiber tracts.
//!
//! The coupling layer bridges the worm body and the neural brain using
//! fibertract-rs. It does NOT assign neurons to channels — instead, it
//! finds the nearest neurons to each bundle's spatial anchor and
//! injects/reads from them by proximity.
//!
//! Early in development, these are undifferentiated neurons. As disc
//! programs fire, they differentiate into proper sensory/motor neurons.
//! The coupling layer doesn't care — it always uses spatial proximity.

use fibertract::{FiberBundle, FiberTractKind, LimbProfile, ReceptorMode};
use fibertract::adapt::{AdaptationConfig, adapt_bundle};
use fibertract::profile::TractSpec;
use neuropool::{Signal, SpatialNeuron, SpatialRuntime};

use crate::body::{
    SensorySnapshot, MotorCommand,
    SEGMENT_COUNT, CHEMO_CHANNELS, DISTRESS_CHANNELS,
    TOUCH_CHANNELS_PER_SEGMENT, PROPRIO_CHANNELS_PER_SEGMENT,
    MOTOR_CHANNELS_PER_SEGMENT,
};

/// Spatial anchor for a bundle — where in brain space this interface lives.
#[derive(Clone, Debug)]
pub struct BundleAnchor {
    /// Name (matches bundle name).
    pub name: String,
    /// 3D position in brain space.
    pub position: [f32; 3],
    /// How far to search for neurons.
    pub search_radius: f32,
}

/// A coupled bundle: fiber tract bundle + spatial anchor + neuron index cache.
#[derive(Clone, Debug)]
pub struct CoupledBundle {
    pub bundle: FiberBundle,
    pub anchor: BundleAnchor,
    /// Cached nearest neuron indices (rebuilt periodically).
    /// One per tract channel.
    pub sensory_neuron_indices: Vec<u32>,
    pub motor_neuron_indices: Vec<u32>,
}

/// The complete coupling state between body and brain.
pub struct CouplingState {
    /// Head bundle: chemosensory + head touch/proprio + head motor.
    pub head: CoupledBundle,
    /// Per-segment bundles: touch + proprio + motor.
    pub segments: Vec<CoupledBundle>,
    /// RNG seed for tract transmission.
    rng_seed: u64,
    /// Fiber tract adaptation config.
    pub adapt_config: AdaptationConfig,
}

/// Sensory scale factor: raw body i32 → neural current injection.
const SENSORY_INJECT_SCALE: f32 = 3.0;


impl CouplingState {
    /// Build coupling state for a worm.
    ///
    /// Creates fiber bundles with spatial anchors matching the body layout.
    pub fn new() -> Self {
        let head = CoupledBundle {
            bundle: head_bundle(),
            anchor: BundleAnchor {
                name: "head".into(),
                position: [0.5, 0.0, 0.5],
                search_radius: 3.0,
            },
            sensory_neuron_indices: Vec::new(),
            motor_neuron_indices: Vec::new(),
        };

        let segments = (0..SEGMENT_COUNT)
            .map(|i| CoupledBundle {
                bundle: segment_bundle(i),
                anchor: BundleAnchor {
                    name: format!("segment_{i}"),
                    position: [-(i as f32), 0.0, 0.5],
                    search_radius: 2.5,
                },
                sensory_neuron_indices: Vec::new(),
                motor_neuron_indices: Vec::new(),
            })
            .collect();

        Self {
            head,
            segments,
            rng_seed: 42,
            adapt_config: AdaptationConfig::default(),
        }
    }

    /// Rebuild neuron index caches based on spatial proximity.
    ///
    /// Call this periodically (or after migration/differentiation) to
    /// update which neurons serve each interface channel.
    ///
    /// Head channels get stable mapping: chemo neurons are sorted by
    /// embodied position (left/right/dorsal/center) rather than arbitrary
    /// proximity order. This is labeled-line identity through embodiment.
    pub fn rebuild_caches(&mut self, neurons: &[SpatialNeuron]) {
        rebuild_bundle_cache(&mut self.head, neurons);
        stabilize_head_channels(&mut self.head, neurons);
        for seg in &mut self.segments {
            rebuild_bundle_cache(seg, neurons);
        }
    }

    /// Transmit sensory input from body to brain.
    ///
    /// Body sensory snapshot → fiber tracts → neural current injection.
    pub fn inject_sensory(
        &mut self,
        snapshot: &SensorySnapshot,
        runtime: &mut SpatialRuntime,
    ) {
        self.rng_seed = self.rng_seed.wrapping_add(1);

        // Head: chemosensory (4 channels) + metabolic distress (1 channel)
        let mut head_sensory: Vec<i32> = snapshot.chemosensory.to_vec();
        head_sensory.push(snapshot.distress);
        inject_sensory_bundle(&mut self.head, &head_sensory, runtime, self.rng_seed);

        // Per-segment: touch (4 channels) + proprioception (2 channels)
        for (seg_idx, coupled) in self.segments.iter_mut().enumerate() {
            let mut seg_sensory = Vec::with_capacity(
                TOUCH_CHANNELS_PER_SEGMENT + PROPRIO_CHANNELS_PER_SEGMENT,
            );
            seg_sensory.extend_from_slice(&snapshot.touch[seg_idx]);
            seg_sensory.extend_from_slice(&snapshot.proprioception[seg_idx]);

            inject_sensory_bundle(coupled, &seg_sensory, runtime, self.rng_seed.wrapping_add(seg_idx as u64));
        }
    }

    /// Read motor output from brain and build motor commands.
    ///
    /// Nearest motor neurons → fire state → fiber tracts → muscle activation.
    pub fn read_motor(
        &mut self,
        runtime: &SpatialRuntime,
    ) -> MotorCommand {
        let mut cmd = MotorCommand::default();

        // Each segment bundle has motor tracts
        for (seg_idx, coupled) in self.segments.iter_mut().enumerate() {
            let motor_signals = read_motor_bundle(coupled, runtime, self.rng_seed.wrapping_add(100 + seg_idx as u64));

            // Map motor signals to muscle activations
            // 4 motor channels per segment: dorsal, ventral, left, right
            for (ch, sig) in motor_signals.iter().enumerate() {
                let activation = if sig.polarity > 0 {
                    sig.magnitude as f32 / 255.0
                } else {
                    0.0
                };

                match ch {
                    0 => cmd.dorsal[seg_idx] = activation,
                    1 => cmd.ventral[seg_idx] = activation,
                    2 => cmd.left[seg_idx] = activation,
                    3 => cmd.right[seg_idx] = activation,
                    _ => {}
                }
            }
        }

        cmd
    }

    /// Get active sensory channels (which channels had nonzero input).
    pub fn active_sensory_channels(&self, snapshot: &SensorySnapshot) -> Vec<u16> {
        let mut channels = Vec::new();
        let mut ch = 0u16;

        // Chemo
        for &v in &snapshot.chemosensory {
            if v.abs() > 10 { channels.push(ch); }
            ch += 1;
        }

        // Distress
        if snapshot.distress.abs() > 10 { channels.push(ch); }
        ch += 1;

        // Touch + proprio per segment
        for seg_idx in 0..SEGMENT_COUNT {
            for &v in &snapshot.touch[seg_idx] {
                if v > 10 { channels.push(ch); }
                ch += 1;
            }
            for &v in &snapshot.proprioception[seg_idx] {
                if v.abs() > 10 { channels.push(ch); }
                ch += 1;
            }
        }

        channels
    }

    /// Get active motor channels (which neurons fired).
    pub fn active_motor_channels(&self, runtime: &SpatialRuntime) -> Vec<u16> {
        let motors = runtime.read_motors();
        motors.iter().map(|&(ch, _)| ch).collect()
    }

    /// Adapt all fiber tracts based on recent activity.
    ///
    /// Call once per tick. Tracts that receive input stimulation will
    /// myelinate and sensitize. Idle tracts demyelinate and desensitize.
    pub fn adapt(&mut self) {
        adapt_bundle(&mut self.head.bundle, &self.adapt_config);
        for seg in &mut self.segments {
            adapt_bundle(&mut seg.bundle, &self.adapt_config);
        }
    }
}

/// Rebuild the neuron index cache for a coupled bundle.
///
/// After differentiation, prefers neurons with matching interfaces:
/// - Sensory-typed neurons for sensory channels
/// - Motor-typed neurons for motor channels
///
/// Falls back to proximity-only assignment when not enough typed neurons
/// exist (e.g., pre-differentiation).
fn rebuild_bundle_cache(coupled: &mut CoupledBundle, neurons: &[SpatialNeuron]) {
    let anchor = coupled.anchor.position;
    let radius = coupled.anchor.search_radius;

    // Find nearest neurons within radius, sorted by distance
    let mut candidates: Vec<(u32, f32)> = neurons
        .iter()
        .enumerate()
        .filter_map(|(idx, n)| {
            let d = distance_3d(n.soma.position, anchor);
            if d < radius {
                Some((idx as u32, d))
            } else {
                None
            }
        })
        .collect();

    candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

    // Count sensory and motor tract channels
    let sensory_count: usize = coupled
        .bundle
        .sensory_tracts()
        .map(|t| t.dim)
        .sum();
    let motor_count: usize = coupled
        .bundle
        .motor_tracts()
        .map(|t| t.dim)
        .sum();

    // Assign typed neurons first, fall back to any nearby neuron
    coupled.sensory_neuron_indices = assign_neurons_prefer_type(
        &candidates, neurons, sensory_count, NeuronRole::Sensory,
    );
    coupled.motor_neuron_indices = assign_neurons_prefer_type(
        &candidates, neurons, motor_count, NeuronRole::Motor,
    );
}

#[derive(Clone, Copy)]
enum NeuronRole { Sensory, Motor }

/// Assign neurons to channels, preferring neurons with matching interface type.
///
/// 1. First fills slots with typed neurons (sensory or motor) sorted by distance.
/// 2. Remaining slots filled by any nearby neuron (proximity fallback).
/// 3. Reuses neurons if not enough candidates.
fn assign_neurons_prefer_type(
    candidates: &[(u32, f32)],
    neurons: &[SpatialNeuron],
    count: usize,
    role: NeuronRole,
) -> Vec<u32> {
    if candidates.is_empty() || count == 0 {
        return Vec::new();
    }

    // Partition into typed (matching interface) and untyped
    let mut typed: Vec<u32> = Vec::new();
    let mut untyped: Vec<u32> = Vec::new();

    for &(idx, _) in candidates {
        let n = &neurons[idx as usize];
        let matches = match role {
            NeuronRole::Sensory => n.nuclei.is_sensory(),
            NeuronRole::Motor => n.nuclei.is_motor(),
        };
        if matches {
            typed.push(idx);
        } else {
            untyped.push(idx);
        }
    }

    // Build assignment: typed first, then untyped fallback
    let mut indices = Vec::with_capacity(count);
    let mut all = typed;
    all.extend(untyped);

    for i in 0..count {
        indices.push(all[i % all.len()]);
    }
    indices
}

/// Stabilize head chemosensory channel mapping by embodied position.
///
/// Sorts the first CHEMO_CHANNELS entries of the head's sensory_neuron_indices
/// so that channels map to neurons by their physical position in the head:
/// - Channel 0 (left gradient)   → neuron with most negative y
/// - Channel 1 (right gradient)  → neuron with most positive y
/// - Channel 2 (dorsal gradient) → neuron with most positive z
/// - Channel 3 (concentration)   → remaining neuron
///
/// This is labeled-line identity through embodiment, not labels. The neuron
/// that is physically leftmost gets the left-gradient signal because it IS
/// on the left.
fn stabilize_head_channels(coupled: &mut CoupledBundle, neurons: &[SpatialNeuron]) {
    if coupled.sensory_neuron_indices.len() < CHEMO_CHANNELS {
        return;
    }

    // Work with the first CHEMO_CHANNELS indices (chemo tract comes first)
    let mut chemo: Vec<u32> = coupled.sensory_neuron_indices[..CHEMO_CHANNELS].to_vec();

    // Deduplicate — if same neuron appears multiple times, we can't sort meaningfully
    chemo.sort();
    chemo.dedup();
    if chemo.len() < CHEMO_CHANNELS {
        return; // not enough distinct neurons to sort
    }

    // Score each neuron on each axis
    let positions: Vec<[f32; 3]> = chemo.iter()
        .map(|&idx| neurons[idx as usize].soma.position)
        .collect();

    let mut assigned = [0u32; 4]; // [left, right, dorsal, concentration]
    let mut used = [false; 4]; // track which of our chemo[] entries are used

    // Channel 0 — left gradient: most negative y
    if let Some(i) = (0..chemo.len())
        .filter(|i| !used[*i])
        .min_by(|&a, &b| positions[a][1].partial_cmp(&positions[b][1]).unwrap())
    {
        assigned[0] = chemo[i]; used[i] = true;
    }

    // Channel 1 — right gradient: most positive y
    if let Some(i) = (0..chemo.len())
        .filter(|i| !used[*i])
        .max_by(|&a, &b| positions[a][1].partial_cmp(&positions[b][1]).unwrap())
    {
        assigned[1] = chemo[i]; used[i] = true;
    }

    // Channel 2 — dorsal gradient: most positive z
    if let Some(i) = (0..chemo.len())
        .filter(|i| !used[*i])
        .max_by(|&a, &b| positions[a][2].partial_cmp(&positions[b][2]).unwrap())
    {
        assigned[2] = chemo[i]; used[i] = true;
    }

    // Channel 3 — concentration: whatever's left
    if let Some(i) = (0..chemo.len()).find(|i| !used[*i]) {
        assigned[3] = chemo[i];
    }

    // Write back stable mapping
    for ch in 0..CHEMO_CHANNELS {
        coupled.sensory_neuron_indices[ch] = assigned[ch];
    }
}

/// Inject sensory signals from a bundle into the brain.
fn inject_sensory_bundle(
    coupled: &mut CoupledBundle,
    raw_sensory: &[i32],
    runtime: &mut SpatialRuntime,
    rng_seed: u64,
) {
    // Transmit through fiber tracts
    let mut tract_idx = 0;
    let mut channel_offset = 0;

    for tract in coupled.bundle.tracts.iter_mut() {
        if !tract.kind.is_afferent() {
            continue;
        }

        let end = (channel_offset + tract.dim).min(raw_sensory.len());
        if channel_offset < end {
            tract.transmit_sensory(&raw_sensory[channel_offset..end], rng_seed.wrapping_add(tract_idx));
        }

        // Inject shaped sensory signals into nearest neurons
        for (ch, &shaped_val) in tract.sensory_signals.iter().enumerate() {
            let neuron_ch = channel_offset + ch;
            if neuron_ch < coupled.sensory_neuron_indices.len() {
                let neuron_idx = coupled.sensory_neuron_indices[neuron_ch];
                if shaped_val != 0 {
                    let current = (shaped_val as f32 * SENSORY_INJECT_SCALE).clamp(-32000.0, 32000.0) as i16;
                    runtime.inject(neuron_idx, current);
                }
            }
        }

        channel_offset = end;
        tract_idx += 1;
    }
}

/// Read motor signals from the brain via a bundle.
fn read_motor_bundle(
    coupled: &mut CoupledBundle,
    runtime: &SpatialRuntime,
    rng_seed: u64,
) -> Vec<Signal> {
    let mut all_motor_signals = Vec::new();
    let mut tract_idx = 0;

    for tract in coupled.bundle.tracts.iter_mut() {
        if !tract.kind.is_efferent() {
            continue;
        }

        // Build input signals from nearest neurons' traces
        let mut input_signals = vec![Signal::default(); tract.dim];
        for (ch, sig) in input_signals.iter_mut().enumerate() {
            let neuron_ch = all_motor_signals.len() + ch;
            if neuron_ch < coupled.motor_neuron_indices.len() {
                let neuron_idx = coupled.motor_neuron_indices[neuron_ch] as usize;
                if neuron_idx < runtime.cascade.neurons.len() {
                    let neuron = &runtime.cascade.neurons[neuron_idx];
                    // Use trace as motor signal strength
                    if neuron.trace > 0 {
                        *sig = Signal {
                            polarity: 1,
                            magnitude: (neuron.trace as u8).min(255),
                        };
                    }
                }
            }
        }

        // Transmit through fiber tract
        tract.transmit_motor(&input_signals, rng_seed.wrapping_add(tract_idx));

        all_motor_signals.extend_from_slice(&tract.motor_signals);
        tract_idx += 1;
    }

    all_motor_signals
}

// =========================================================================
// Worm body fiber tract profiles
// =========================================================================

/// Head bundle: chemosensory + metabolic distress tracts.
///
/// Chemosensory receptors are tonic — the worm constantly senses the
/// chemical gradient. Metabolic distress is a separate interoceptive
/// channel carrying internal energy state. All properties start at
/// biological defaults and adapt through use.
fn head_bundle() -> FiberBundle {
    let profile = LimbProfile {
        name: "head".into(),
        tracts: vec![
            // Chemosensory: external gradient detection
            TractSpec {
                kind: FiberTractKind::Interoceptive,
                dim: CHEMO_CHANNELS,
                receptor_mode: Some(ReceptorMode::Tonic),
                ..TractSpec::new(FiberTractKind::Interoceptive, CHEMO_CHANNELS)
            },
            // Metabolic distress: internal energy state
            TractSpec {
                kind: FiberTractKind::Interoceptive,
                dim: DISTRESS_CHANNELS,
                receptor_mode: Some(ReceptorMode::Tonic),
                ..TractSpec::new(FiberTractKind::Interoceptive, DISTRESS_CHANNELS)
            },
        ],
    };
    profile.build()
}

/// Per-segment bundle: touch + proprioception + motor.
///
/// Body wall receptors are tonic (sustained level reporting).
/// Motor tracts start at defaults — adaptation will myelinate active
/// tracts, sensitize stimulated ones, and let the system self-tune.
fn segment_bundle(segment_idx: usize) -> FiberBundle {
    let name = format!("segment_{segment_idx}");
    let profile = LimbProfile {
        name,
        tracts: vec![
            // Touch: tonic mechanoreceptors (body wall)
            TractSpec {
                kind: FiberTractKind::Mechanoreceptive,
                dim: TOUCH_CHANNELS_PER_SEGMENT,
                receptor_mode: Some(ReceptorMode::Tonic),
                ..TractSpec::new(FiberTractKind::Mechanoreceptive, TOUCH_CHANNELS_PER_SEGMENT)
            },
            // Proprioception: tonic muscle spindles
            TractSpec {
                kind: FiberTractKind::Proprioceptive,
                dim: PROPRIO_CHANNELS_PER_SEGMENT,
                receptor_mode: Some(ReceptorMode::Tonic),
                ..TractSpec::new(FiberTractKind::Proprioceptive, PROPRIO_CHANNELS_PER_SEGMENT)
            },
            // Motor: starts at defaults, adapts through stimulation
            TractSpec::new(FiberTractKind::MotorSkeletal, MOTOR_CHANNELS_PER_SEGMENT),
        ],
    };
    profile.build()
}

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()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::body::{Environment, WormBody};

    #[test]
    fn coupling_creation() {
        let coupling = CouplingState::new();
        assert_eq!(coupling.segments.len(), SEGMENT_COUNT);
    }

    #[test]
    fn head_bundle_has_sensory() {
        let bundle = head_bundle();
        assert!(bundle.sensory_tracts().count() > 0);
    }

    #[test]
    fn segment_bundle_has_motor_and_sensory() {
        let bundle = segment_bundle(0);
        assert!(bundle.sensory_tracts().count() > 0);
        assert!(bundle.motor_tracts().count() > 0);
    }

    #[test]
    fn all_segment_bundles_consistent() {
        for i in 0..SEGMENT_COUNT {
            let b = segment_bundle(i);
            assert_eq!(b.name, format!("segment_{i}"));
            assert!(b.tract_count() >= 3);
        }
    }

    #[test]
    fn sensory_snapshot_maps_to_channels() {
        let body = WormBody::new(Environment::default());
        let snap = body.sense();

        let coupling = CouplingState::new();
        let active = coupling.active_sensory_channels(&snap);
        // Ground contact should generate ventral touch on all segments
        assert!(!active.is_empty(), "ground contact should produce sensory activity");
    }
}