ruvector-nervous-system 2.0.6

//! E-prop (Eligibility Propagation) online learning algorithm.
//!
//! Based on Bellec et al. 2020: "A solution to the learning dilemma for recurrent
//! networks of spiking neurons"
//!
//! ## Key Features
//!
//! - **O(1) Memory**: Only 12 bytes per synapse (weight + 2 traces)
//! - **No BPTT**: No need for backpropagation through time
//! - **Long Credit Assignment**: Handles 1000+ millisecond temporal windows
//! - **Three-Factor Rule**: Δw = η × eligibility_trace × learning_signal
//!
//! ## Algorithm
//!
//! 1. **Eligibility Traces**: Exponentially decaying traces capture pre-post correlations
//! 2. **Surrogate Gradients**: Pseudo-derivatives enable gradient-based learning in SNNs
//! 3. **Learning Signals**: Broadcast error signals from output layer
//! 4. **Local Updates**: All computations are local to each synapse

use rand_distr::{Distribution, Normal};
use serde::{Deserialize, Serialize};

/// E-prop synapse with eligibility traces for online learning.
///
/// Memory footprint: 12 bytes (3 × f32)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpropSynapse {
    /// Synaptic weight
    pub weight: f32,

    /// Eligibility trace (fast component)
    pub eligibility_trace: f32,

    /// Filtered eligibility trace (slow component for stability)
    pub filtered_trace: f32,

    /// Time constant for eligibility trace decay (ms)
    pub tau_e: f32,

    /// Time constant for slow trace filter (ms)
    pub tau_slow: f32,
}

impl EpropSynapse {
    /// Create new synapse with random initialization.
    ///
    /// # Arguments
    ///
    /// * `initial_weight` - Initial synaptic weight
    /// * `tau_e` - Eligibility trace time constant (10-1000 ms)
    ///
    /// # Example
    ///
    /// ```
    /// use ruvector_nervous_system::plasticity::eprop::EpropSynapse;
    ///
    /// let synapse = EpropSynapse::new(0.1, 20.0);
    /// ```
    pub fn new(initial_weight: f32, tau_e: f32) -> Self {
        Self {
            weight: initial_weight,
            eligibility_trace: 0.0,
            filtered_trace: 0.0,
            tau_e,
            tau_slow: tau_e * 2.0, // Slow filter is 2x the fast trace
        }
    }

    /// Update synapse using three-factor learning rule.
    ///
    /// # Arguments
    ///
    /// * `pre_spike` - Did presynaptic neuron spike?
    /// * `pseudo_derivative` - Surrogate gradient from postsynaptic neuron
    /// * `learning_signal` - Error signal from output layer
    /// * `dt` - Time step (ms)
    /// * `lr` - Learning rate
    ///
    /// # Three-Factor Rule
    ///
    /// ```text
    /// Δw = η × e × L
    /// where:
    ///   η = learning rate
    ///   e = eligibility trace
    ///   L = learning signal
    /// ```
    pub fn update(
        &mut self,
        pre_spike: bool,
        pseudo_derivative: f32,
        learning_signal: f32,
        dt: f32,
        lr: f32,
    ) {
        // Decay eligibility traces exponentially
        let decay_fast = (-dt / self.tau_e).exp();
        let decay_slow = (-dt / self.tau_slow).exp();

        self.eligibility_trace *= decay_fast;
        self.filtered_trace *= decay_slow;

        // Accumulate trace on presynaptic spike
        if pre_spike {
            let trace_increment = pseudo_derivative;
            self.eligibility_trace += trace_increment;
            self.filtered_trace += trace_increment;
        }

        // Three-factor weight update
        // Use filtered trace for stability
        let weight_delta = lr * self.filtered_trace * learning_signal;
        self.weight += weight_delta;

        // Optional: weight clipping for stability
        self.weight = self.weight.clamp(-10.0, 10.0);
    }

    /// Reset eligibility traces (e.g., between trials).
    pub fn reset_traces(&mut self) {
        self.eligibility_trace = 0.0;
        self.filtered_trace = 0.0;
    }
}

/// Leaky Integrate-and-Fire (LIF) neuron for E-prop.
///
/// Implements surrogate gradient (pseudo-derivative) for backprop compatibility.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpropLIF {
    /// Membrane potential (mV)
    pub membrane: f32,

    /// Spike threshold (mV)
    pub threshold: f32,

    /// Membrane time constant (ms)
    pub tau_mem: f32,

    /// Refractory period counter (ms)
    pub refractory: u32,

    /// Refractory period duration (ms)
    pub refractory_period: u32,

    /// Resting potential (mV)
    pub v_rest: f32,

    /// Reset potential (mV)
    pub v_reset: f32,
}

impl EpropLIF {
    /// Create new LIF neuron with default parameters.
    ///
    /// # Example
    ///
    /// ```
    /// use ruvector_nervous_system::plasticity::eprop::EpropLIF;
    ///
    /// let neuron = EpropLIF::new(-70.0, -55.0, 20.0);
    /// ```
    pub fn new(v_rest: f32, threshold: f32, tau_mem: f32) -> Self {
        Self {
            membrane: v_rest,
            threshold,
            tau_mem,
            refractory: 0,
            refractory_period: 2, // 2 ms refractory period
            v_rest,
            v_reset: v_rest,
        }
    }

    /// Step neuron forward in time.
    ///
    /// Returns `(spike, pseudo_derivative)` where:
    /// - `spike`: true if neuron spiked this timestep
    /// - `pseudo_derivative`: surrogate gradient for learning
    ///
    /// # Surrogate Gradient
    ///
    /// Uses fast sigmoid approximation:
    /// ```text
    /// σ'(V) = max(0, 1 - |V - θ|)
    /// ```
    pub fn step(&mut self, input: f32, dt: f32) -> (bool, f32) {
        let mut spike = false;
        let mut pseudo_derivative = 0.0;

        // Handle refractory period
        if self.refractory > 0 {
            self.refractory -= 1;
            self.membrane = self.v_reset;
            return (false, 0.0);
        }

        // Leaky integration
        let decay = (-dt / self.tau_mem).exp();
        self.membrane = self.membrane * decay + input * (1.0 - decay);

        // Compute pseudo-derivative (before spike check)
        // Fast sigmoid: max(0, 1 - |V - threshold|)
        let distance = (self.membrane - self.threshold).abs();
        pseudo_derivative = (1.0 - distance).max(0.0);

        // Spike generation
        if self.membrane >= self.threshold {
            spike = true;
            self.membrane = self.v_reset;
            self.refractory = self.refractory_period;
        }

        (spike, pseudo_derivative)
    }

    /// Reset neuron state.
    pub fn reset(&mut self) {
        self.membrane = self.v_rest;
        self.refractory = 0;
    }
}

/// Learning signal generation strategies.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LearningSignal {
    /// Symmetric e-prop: direct error propagation
    Symmetric(f32),

    /// Random feedback alignment
    Random { feedback: Vec<f32> },

    /// Adaptive learning signal with buffer
    Adaptive { buffer: Vec<f32> },
}

impl LearningSignal {
    /// Compute learning signal for a neuron.
    pub fn compute(&self, neuron_idx: usize, error: f32) -> f32 {
        match self {
            LearningSignal::Symmetric(scale) => error * scale,
            LearningSignal::Random { feedback } => {
                if neuron_idx < feedback.len() {
                    error * feedback[neuron_idx]
                } else {
                    0.0
                }
            }
            LearningSignal::Adaptive { buffer } => {
                if neuron_idx < buffer.len() {
                    error * buffer[neuron_idx]
                } else {
                    0.0
                }
            }
        }
    }
}

/// E-prop recurrent neural network.
///
/// Three-layer architecture: Input → Recurrent Hidden → Readout
///
/// # Performance
///
/// - Per-synapse memory: 12 bytes
/// - Update time: <1 ms for 1000 neurons, 100k synapses
/// - Credit assignment: 1000+ ms temporal windows
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpropNetwork {
    /// Input size
    pub input_size: usize,

    /// Hidden layer size
    pub hidden_size: usize,

    /// Output size
    pub output_size: usize,

    /// Hidden layer neurons
    pub neurons: Vec<EpropLIF>,

    /// Input → Hidden synapses (input_size × hidden_size)
    pub input_synapses: Vec<Vec<EpropSynapse>>,

    /// Recurrent Hidden → Hidden synapses (hidden_size × hidden_size)
    pub recurrent_synapses: Vec<Vec<EpropSynapse>>,

    /// Hidden → Output readout weights (hidden_size × output_size)
    pub readout: Vec<Vec<f32>>,

    /// Learning signal strategy
    pub learning_signal: LearningSignal,

    /// Hidden layer spike buffer
    spike_buffer: Vec<bool>,

    /// Hidden layer pseudo-derivatives
    pseudo_derivatives: Vec<f32>,
}

impl EpropNetwork {
    /// Create new E-prop network.
    ///
    /// # Arguments
    ///
    /// * `input_size` - Number of input neurons
    /// * `hidden_size` - Number of hidden recurrent neurons
    /// * `output_size` - Number of output neurons
    ///
    /// # Example
    ///
    /// ```
    /// use ruvector_nervous_system::plasticity::eprop::EpropNetwork;
    ///
    /// // Create network: 28×28 input, 256 hidden, 10 output
    /// let network = EpropNetwork::new(784, 256, 10);
    /// ```
    pub fn new(input_size: usize, hidden_size: usize, output_size: usize) -> Self {
        let mut rng = rand::thread_rng();

        // Initialize hidden neurons
        let neurons = (0..hidden_size)
            .map(|_| EpropLIF::new(-70.0, -55.0, 20.0))
            .collect();

        // Initialize input synapses with He initialization
        let input_scale = (2.0 / input_size as f32).sqrt();
        let normal = Normal::new(0.0, input_scale as f64).unwrap();
        let input_synapses = (0..input_size)
            .map(|_| {
                (0..hidden_size)
                    .map(|_| {
                        let weight = normal.sample(&mut rng) as f32;
                        EpropSynapse::new(weight, 20.0)
                    })
                    .collect()
            })
            .collect();

        // Initialize recurrent synapses (sparser initialization)
        let recurrent_scale = (1.0 / hidden_size as f32).sqrt();
        let recurrent_normal = Normal::new(0.0, recurrent_scale as f64).unwrap();
        let recurrent_synapses = (0..hidden_size)
            .map(|_| {
                (0..hidden_size)
                    .map(|_| {
                        let weight = recurrent_normal.sample(&mut rng) as f32;
                        EpropSynapse::new(weight, 20.0)
                    })
                    .collect()
            })
            .collect();

        // Initialize readout layer
        let readout_scale = (1.0 / hidden_size as f32).sqrt();
        let readout_normal = Normal::new(0.0, readout_scale as f64).unwrap();
        let readout = (0..hidden_size)
            .map(|_| {
                (0..output_size)
                    .map(|_| readout_normal.sample(&mut rng) as f32)
                    .collect()
            })
            .collect();

        Self {
            input_size,
            hidden_size,
            output_size,
            neurons,
            input_synapses,
            recurrent_synapses,
            readout,
            learning_signal: LearningSignal::Symmetric(1.0),
            spike_buffer: vec![false; hidden_size],
            pseudo_derivatives: vec![0.0; hidden_size],
        }
    }

    /// Forward pass through network.
    ///
    /// # Arguments
    ///
    /// * `input` - Input spike train (0 or 1 for each input neuron)
    /// * `dt` - Time step in milliseconds
    ///
    /// # Returns
    ///
    /// Output activations (readout layer)
    pub fn forward(&mut self, input: &[f32], dt: f32) -> Vec<f32> {
        assert_eq!(input.len(), self.input_size, "Input size mismatch");

        // Compute input currents
        let mut currents = vec![0.0; self.hidden_size];

        // Input → Hidden
        for (i, &inp) in input.iter().enumerate() {
            if inp > 0.5 {
                // Input spike
                for (j, synapse) in self.input_synapses[i].iter().enumerate() {
                    currents[j] += synapse.weight;
                }
            }
        }

        // Recurrent Hidden → Hidden (using previous spike buffer)
        for (i, &spike) in self.spike_buffer.iter().enumerate() {
            if spike {
                for (j, synapse) in self.recurrent_synapses[i].iter().enumerate() {
                    currents[j] += synapse.weight;
                }
            }
        }

        // Update neurons
        for (i, neuron) in self.neurons.iter_mut().enumerate() {
            let (spike, pseudo_deriv) = neuron.step(currents[i], dt);
            self.spike_buffer[i] = spike;
            self.pseudo_derivatives[i] = pseudo_deriv;
        }

        // Readout layer (linear)
        let mut output = vec![0.0; self.output_size];
        for (i, &spike) in self.spike_buffer.iter().enumerate() {
            if spike {
                for (j, weight) in self.readout[i].iter().enumerate() {
                    output[j] += weight;
                }
            }
        }

        output
    }

    /// Backward pass: update eligibility traces and weights.
    ///
    /// # Arguments
    ///
    /// * `error` - Error signal from output layer (target - prediction)
    /// * `learning_rate` - Learning rate
    /// * `dt` - Time step in milliseconds
    pub fn backward(&mut self, error: &[f32], learning_rate: f32, dt: f32) {
        assert_eq!(error.len(), self.output_size, "Error size mismatch");

        // Compute learning signals for each hidden neuron
        let mut learning_signals = vec![0.0; self.hidden_size];

        // Backpropagate error through readout layer
        for i in 0..self.hidden_size {
            let mut signal = 0.0;
            for j in 0..self.output_size {
                signal += error[j] * self.readout[i][j];
            }
            learning_signals[i] = self.learning_signal.compute(i, signal);
        }

        // Update input synapses
        for i in 0..self.input_size {
            for j in 0..self.hidden_size {
                // Check if input spiked (simplified: use input value)
                let pre_spike = false; // Will be set by caller tracking input history
                self.input_synapses[i][j].update(
                    pre_spike,
                    self.pseudo_derivatives[j],
                    learning_signals[j],
                    dt,
                    learning_rate,
                );
            }
        }

        // Update recurrent synapses
        for i in 0..self.hidden_size {
            for j in 0..self.hidden_size {
                let pre_spike = self.spike_buffer[i];
                self.recurrent_synapses[i][j].update(
                    pre_spike,
                    self.pseudo_derivatives[j],
                    learning_signals[j],
                    dt,
                    learning_rate,
                );
            }
        }

        // Update readout weights (simple gradient descent)
        for i in 0..self.hidden_size {
            if self.spike_buffer[i] {
                for j in 0..self.output_size {
                    self.readout[i][j] += learning_rate * error[j];
                }
            }
        }
    }

    /// Single online learning step (forward + backward).
    ///
    /// # Arguments
    ///
    /// * `input` - Input spike train
    /// * `target` - Target output
    /// * `dt` - Time step (ms)
    /// * `lr` - Learning rate
    ///
    /// # Example
    ///
    /// ```
    /// use ruvector_nervous_system::plasticity::eprop::EpropNetwork;
    ///
    /// let mut network = EpropNetwork::new(10, 100, 2);
    ///
    /// // Training loop
    /// for _ in 0..1000 {
    ///     let input = vec![0.0; 10];
    ///     let target = vec![1.0, 0.0];
    ///     network.online_step(&input, &target, 1.0, 0.001);
    /// }
    /// ```
    pub fn online_step(&mut self, input: &[f32], target: &[f32], dt: f32, lr: f32) {
        let output = self.forward(input, dt);

        // Compute error
        let error: Vec<f32> = target
            .iter()
            .zip(output.iter())
            .map(|(t, o)| t - o)
            .collect();

        self.backward(&error, lr, dt);
    }

    /// Reset network state (neurons and eligibility traces).
    pub fn reset(&mut self) {
        for neuron in &mut self.neurons {
            neuron.reset();
        }

        for synapses in &mut self.input_synapses {
            for synapse in synapses {
                synapse.reset_traces();
            }
        }

        for synapses in &mut self.recurrent_synapses {
            for synapse in synapses {
                synapse.reset_traces();
            }
        }

        self.spike_buffer.fill(false);
        self.pseudo_derivatives.fill(0.0);
    }

    /// Get total number of synapses.
    pub fn num_synapses(&self) -> usize {
        let input_synapses = self.input_size * self.hidden_size;
        let recurrent_synapses = self.hidden_size * self.hidden_size;
        let readout_synapses = self.hidden_size * self.output_size;
        input_synapses + recurrent_synapses + readout_synapses
    }

    /// Estimate memory footprint in bytes.
    pub fn memory_footprint(&self) -> usize {
        let synapse_size = std::mem::size_of::<EpropSynapse>();
        let neuron_size = std::mem::size_of::<EpropLIF>();
        let readout_size = std::mem::size_of::<f32>();

        let input_mem = self.input_size * self.hidden_size * synapse_size;
        let recurrent_mem = self.hidden_size * self.hidden_size * synapse_size;
        let readout_mem = self.hidden_size * self.output_size * readout_size;
        let neuron_mem = self.hidden_size * neuron_size;

        input_mem + recurrent_mem + readout_mem + neuron_mem
    }
}

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

    #[test]
    fn test_synapse_creation() {
        let synapse = EpropSynapse::new(0.5, 20.0);
        assert_eq!(synapse.weight, 0.5);
        assert_eq!(synapse.eligibility_trace, 0.0);
        assert_eq!(synapse.tau_e, 20.0);
    }

    #[test]
    fn test_trace_decay() {
        let mut synapse = EpropSynapse::new(0.5, 20.0);
        synapse.eligibility_trace = 1.0;

        // Decay over 20ms (one time constant)
        synapse.update(false, 0.0, 0.0, 20.0, 0.0);

        // Should decay to ~1/e ≈ 0.368
        assert!((synapse.eligibility_trace - 0.368).abs() < 0.01);
    }

    #[test]
    fn test_lif_spike_generation() {
        let mut neuron = EpropLIF::new(-70.0, -55.0, 20.0);

        // Apply strong input repeatedly to reach threshold
        // With tau=20ms and input=100, need several steps
        for _ in 0..50 {
            let (spike, _) = neuron.step(100.0, 1.0);
            if spike {
                assert_eq!(neuron.membrane, neuron.v_reset);
                return;
            }
        }
        // Should have spiked by now
        panic!("Neuron did not spike with strong sustained input");
    }

    #[test]
    fn test_lif_refractory_period() {
        let mut neuron = EpropLIF::new(-70.0, -55.0, 20.0);

        // First reach threshold and spike
        for _ in 0..50 {
            let (spike, _) = neuron.step(100.0, 1.0);
            if spike {
                break;
            }
        }

        // Try to spike again immediately
        let (spike2, _) = neuron.step(100.0, 1.0);

        // Should not spike (refractory)
        assert!(!spike2, "Should be in refractory period");
    }

    #[test]
    fn test_pseudo_derivative() {
        let mut neuron = EpropLIF::new(-70.0, -55.0, 20.0);

        // Set membrane close to threshold for non-zero pseudo-derivative
        neuron.membrane = -55.5; // Just below threshold

        let (_, pseudo_deriv) = neuron.step(0.0, 1.0);

        // Pseudo-derivative = max(0, 1 - |V - threshold|)
        // With V = -55.5 after decay, distance from -55 should be small
        // The derivative should be >= 0 (may be 0 if distance > 1)
        assert!(pseudo_deriv >= 0.0, "pseudo_deriv={}", pseudo_deriv);
    }

    #[test]
    fn test_network_creation() {
        let network = EpropNetwork::new(10, 100, 2);

        assert_eq!(network.input_size, 10);
        assert_eq!(network.hidden_size, 100);
        assert_eq!(network.output_size, 2);
        assert_eq!(network.neurons.len(), 100);
    }

    #[test]
    fn test_network_forward() {
        let mut network = EpropNetwork::new(10, 50, 2);

        let input = vec![1.0; 10];
        let output = network.forward(&input, 1.0);

        assert_eq!(output.len(), 2);
    }

    #[test]
    fn test_network_memory_footprint() {
        let network = EpropNetwork::new(100, 500, 10);

        let footprint = network.memory_footprint();
        let num_synapses = network.num_synapses();

        // Should be roughly 12 bytes per synapse
        let bytes_per_synapse = footprint / num_synapses;
        assert!(bytes_per_synapse >= 10 && bytes_per_synapse <= 20);
    }

    #[test]
    fn test_online_learning() {
        let mut network = EpropNetwork::new(10, 50, 2);

        let input = vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
        let target = vec![1.0, 0.0];

        // Run several learning steps
        for _ in 0..10 {
            network.online_step(&input, &target, 1.0, 0.01);
        }

        // Network should run without panic
    }

    #[test]
    fn test_network_reset() {
        let mut network = EpropNetwork::new(10, 50, 2);

        // Run forward pass
        let input = vec![1.0; 10];
        network.forward(&input, 1.0);

        // Reset
        network.reset();

        // All neurons should be at rest
        for neuron in &network.neurons {
            assert_eq!(neuron.membrane, neuron.v_rest);
            assert_eq!(neuron.refractory, 0);
        }
    }
}