ruvector-nervous-system 2.0.6

//! Predictive coding layer with residual gating
//!
//! Based on predictive coding theory: only transmit prediction errors (residuals)
//! when they exceed a threshold. This achieves 90-99% bandwidth reduction by
//! suppressing predictable signals.

use std::f32;

/// Predictive layer that learns to predict input and only transmits residuals
#[derive(Debug, Clone)]
pub struct PredictiveLayer {
    /// Current prediction of input
    prediction: Vec<f32>,
    /// Threshold for residual transmission (e.g., 0.1 for 10% change)
    residual_threshold: f32,
    /// Learning rate for prediction updates
    learning_rate: f32,
}

impl PredictiveLayer {
    /// Create a new predictive layer
    ///
    /// # Arguments
    /// * `size` - Dimension of input/prediction vectors
    /// * `threshold` - Residual threshold for transmission (0.0-1.0)
    pub fn new(size: usize, threshold: f32) -> Self {
        Self {
            prediction: vec![0.0; size],
            residual_threshold: threshold,
            learning_rate: 0.1,
        }
    }

    /// Create with custom learning rate
    pub fn with_learning_rate(size: usize, threshold: f32, learning_rate: f32) -> Self {
        Self {
            prediction: vec![0.0; size],
            residual_threshold: threshold,
            learning_rate,
        }
    }

    /// Compute prediction error (residual) between prediction and actual
    ///
    /// # Returns
    /// Vector of residuals (actual - prediction)
    pub fn compute_residual(&self, actual: &[f32]) -> Vec<f32> {
        assert_eq!(actual.len(), self.prediction.len(), "Input size mismatch");

        actual
            .iter()
            .zip(self.prediction.iter())
            .map(|(a, p)| a - p)
            .collect()
    }

    /// Check if residual exceeds threshold and should be transmitted
    ///
    /// Uses RMS (root mean square) of residual as the decision metric
    pub fn should_transmit(&self, actual: &[f32]) -> bool {
        let residual = self.compute_residual(actual);
        let rms = self.residual_rms(&residual);
        rms > self.residual_threshold
    }

    /// Update prediction based on actual input (learning step)
    ///
    /// Uses exponential moving average: prediction = (1-α)*prediction + α*actual
    pub fn update_prediction(&mut self, actual: &[f32], learning_rate: f32) {
        assert_eq!(actual.len(), self.prediction.len(), "Input size mismatch");

        for (pred, &act) in self.prediction.iter_mut().zip(actual.iter()) {
            *pred = (1.0 - learning_rate) * *pred + learning_rate * act;
        }
    }

    /// Update prediction with the layer's default learning rate
    pub fn update(&mut self, actual: &[f32]) {
        self.update_prediction(actual, self.learning_rate);
    }

    /// Perform residual-gated write: only transmit if residual exceeds threshold
    ///
    /// # Returns
    /// * `Some(residual)` if transmission threshold exceeded
    /// * `None` if prediction is good enough (no transmission needed)
    pub fn residual_gated_write(&mut self, actual: &[f32]) -> Option<Vec<f32>> {
        if self.should_transmit(actual) {
            let residual = self.compute_residual(actual);
            self.update(actual);
            Some(residual)
        } else {
            // Update prediction even when not transmitting
            self.update(actual);
            None
        }
    }

    /// Get current prediction (for debugging/analysis)
    pub fn prediction(&self) -> &[f32] {
        &self.prediction
    }

    /// Set residual threshold
    pub fn set_threshold(&mut self, threshold: f32) {
        self.residual_threshold = threshold;
    }

    /// Get residual threshold
    pub fn threshold(&self) -> f32 {
        self.residual_threshold
    }

    /// Compute RMS (root mean square) of residual vector
    fn residual_rms(&self, residual: &[f32]) -> f32 {
        if residual.is_empty() {
            return 0.0;
        }

        let sum_squares: f32 = residual.iter().map(|r| r * r).sum();
        (sum_squares / residual.len() as f32).sqrt()
    }

    /// Get compression ratio (fraction of transmissions)
    ///
    /// Track over a window of attempts
    pub fn compression_stats(&self, attempts: &[bool]) -> f32 {
        if attempts.is_empty() {
            return 0.0;
        }

        let transmissions = attempts.iter().filter(|&&x| x).count();
        transmissions as f32 / attempts.len() as f32
    }
}

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

    #[test]
    fn test_new_predictive_layer() {
        let layer = PredictiveLayer::new(10, 0.1);
        assert_eq!(layer.prediction.len(), 10);
        assert_eq!(layer.residual_threshold, 0.1);
        assert!(layer.prediction.iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_compute_residual() {
        let layer = PredictiveLayer::new(3, 0.1);
        let actual = vec![1.0, 2.0, 3.0];
        let residual = layer.compute_residual(&actual);

        assert_eq!(residual, vec![1.0, 2.0, 3.0]); // prediction is all zeros
    }

    #[test]
    fn test_update_prediction() {
        let mut layer = PredictiveLayer::new(3, 0.1);
        let actual = vec![1.0, 2.0, 3.0];

        layer.update_prediction(&actual, 0.5);

        // prediction = 0.5 * 0.0 + 0.5 * actual
        assert_eq!(layer.prediction, vec![0.5, 1.0, 1.5]);
    }

    #[test]
    fn test_should_transmit() {
        let layer = PredictiveLayer::new(4, 0.5);

        // Small change - should not transmit
        let small_change = vec![0.1, 0.1, 0.1, 0.1];
        assert!(!layer.should_transmit(&small_change));

        // Large change - should transmit
        let large_change = vec![1.0, 1.0, 1.0, 1.0];
        assert!(layer.should_transmit(&large_change));
    }

    #[test]
    fn test_residual_gated_write() {
        let mut layer = PredictiveLayer::new(4, 0.5);

        // Small change - no transmission
        let small_change = vec![0.1, 0.1, 0.1, 0.1];
        let result = layer.residual_gated_write(&small_change);
        assert!(result.is_none());

        // Large change - should transmit residual
        let large_change = vec![1.0, 1.0, 1.0, 1.0];
        let result = layer.residual_gated_write(&large_change);
        assert!(result.is_some());

        let residual = result.unwrap();
        assert!(residual.iter().all(|&r| r.abs() > 0.0));
    }

    #[test]
    fn test_prediction_convergence() {
        let mut layer = PredictiveLayer::with_learning_rate(3, 0.1, 0.2);
        let signal = vec![1.0, 2.0, 3.0];

        // Repeat same signal - prediction should converge
        for _ in 0..50 {
            // More iterations for convergence
            layer.update(&signal);
        }

        // Prediction should be close to signal (relaxed tolerance)
        for (pred, &actual) in layer.prediction.iter().zip(signal.iter()) {
            assert!(
                (pred - actual).abs() < 0.05,
                "Prediction {} did not converge to {}",
                pred,
                actual
            );
        }
    }

    #[test]
    fn test_compression_ratio() {
        let mut layer = PredictiveLayer::new(4, 0.3);
        let mut attempts = Vec::new();

        // Stable signal - should quickly learn and stop transmitting
        let stable_signal = vec![1.0, 1.0, 1.0, 1.0];

        for _ in 0..100 {
            let transmitted = layer.should_transmit(&stable_signal);
            attempts.push(transmitted);
            layer.update(&stable_signal);
        }

        let compression = layer.compression_stats(&attempts);

        // Should transmit less as prediction improves
        // After 100 iterations, compression should be high (low transmission rate)
        assert!(
            compression < 0.5,
            "Compression ratio too low: {}",
            compression
        );
    }

    #[test]
    fn test_residual_rms() {
        let layer = PredictiveLayer::new(4, 0.1);

        // RMS of [1,1,1,1] should be 1.0
        let residual = vec![1.0, 1.0, 1.0, 1.0];
        let rms = layer.residual_rms(&residual);
        assert!((rms - 1.0).abs() < 0.001);

        // RMS of [0,0,0,0] should be 0.0
        let zero_residual = vec![0.0, 0.0, 0.0, 0.0];
        let rms = layer.residual_rms(&zero_residual);
        assert_eq!(rms, 0.0);
    }

    #[test]
    fn test_bandwidth_reduction() {
        let mut layer = PredictiveLayer::with_learning_rate(8, 0.2, 0.3);
        let mut transmission_count = 0;
        let total_attempts = 1000;

        // Slowly varying signal (simulates typical neural activity)
        let mut signal = vec![0.0; 8];

        for i in 0..total_attempts {
            // Add small random perturbation
            let noise = (i as f32 * 0.01).sin() * 0.1;
            signal[0] = 1.0 + noise;

            if layer.residual_gated_write(&signal).is_some() {
                transmission_count += 1;
            }
        }

        let reduction = 1.0 - (transmission_count as f32 / total_attempts as f32);

        // Should achieve at least 50% bandwidth reduction
        assert!(
            reduction > 0.5,
            "Bandwidth reduction too low: {:.1}%",
            reduction * 100.0
        );
    }
}