rustytorch/crates/specialized/rtx-neuro-lsl/src/buffer.rs

//! Ring buffer for continuous streaming data

use std::collections::VecDeque;

use crate::stream::LslSample;

/// Ring buffer for storing streaming samples
pub struct StreamBuffer {
    /// Sample data storage (each Vec is one sample)
    data: VecDeque<Vec<f64>>,
    /// Timestamps for each sample
    timestamps: VecDeque<f64>,
    /// Maximum capacity in samples
    capacity: usize,
    /// Sampling frequency in Hz
    sfreq: f64,
    /// Total samples received (including dropped)
    total_received: usize,
    /// Samples dropped due to buffer overflow
    dropped: usize,
}

impl StreamBuffer {
    /// Create a new buffer with given capacity
    pub fn new(capacity: usize, sfreq: f64) -> Self {
        Self {
            data: VecDeque::with_capacity(capacity),
            timestamps: VecDeque::with_capacity(capacity),
            capacity,
            sfreq,
            total_received: 0,
            dropped: 0,
        }
    }

    /// Create a buffer sized for a given duration
    pub fn for_duration(duration_sec: f64, sfreq: f64) -> Self {
        let capacity = (duration_sec * sfreq).ceil() as usize;
        Self::new(capacity.max(1), sfreq)
    }

    /// Push a new sample into the buffer
    pub fn push_sample(&mut self, sample: Vec<f64>, timestamp: f64) {
        self.total_received += 1;

        // Drop oldest sample if at capacity
        if self.data.len() >= self.capacity {
            self.data.pop_front();
            self.timestamps.pop_front();
            self.dropped += 1;
        }

        self.data.push_back(sample);
        self.timestamps.push_back(timestamp);
    }

    /// Push a sample from LslSample struct
    pub fn push(&mut self, sample: LslSample) {
        self.push_sample(sample.data, sample.timestamp);
    }

    /// Get the most recent samples covering the specified duration
    pub fn get_recent(&self, duration_sec: f64) -> (Vec<Vec<f64>>, Vec<f64>) {
        let n_samples = (duration_sec * self.sfreq).ceil() as usize;
        let n_available = self.data.len();
        let n_to_return = n_samples.min(n_available);

        let start_idx = n_available.saturating_sub(n_to_return);

        let data: Vec<Vec<f64>> = self.data.iter().skip(start_idx).cloned().collect();
        let timestamps: Vec<f64> = self.timestamps.iter().skip(start_idx).copied().collect();

        (data, timestamps)
    }

    /// Get all samples in the buffer
    pub fn get_all(&self) -> (Vec<Vec<f64>>, Vec<f64>) {
        (
            self.data.iter().cloned().collect(),
            self.timestamps.iter().copied().collect(),
        )
    }

    /// Clear the buffer
    pub fn clear(&mut self) {
        self.data.clear();
        self.timestamps.clear();
    }

    /// Get the number of samples currently in the buffer
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Check if buffer is empty
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Get buffer capacity
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Get total samples received
    pub fn total_received(&self) -> usize {
        self.total_received
    }

    /// Get number of dropped samples
    pub fn dropped(&self) -> usize {
        self.dropped
    }

    /// Get the sampling frequency
    pub fn sfreq(&self) -> f64 {
        self.sfreq
    }

    /// Get duration of buffered data in seconds
    pub fn duration(&self) -> f64 {
        if self.data.is_empty() {
            0.0
        } else {
            self.data.len() as f64 / self.sfreq
        }
    }

    /// Get the latest timestamp
    pub fn latest_timestamp(&self) -> Option<f64> {
        self.timestamps.back().copied()
    }

    /// Get the earliest timestamp
    pub fn earliest_timestamp(&self) -> Option<f64> {
        self.timestamps.front().copied()
    }
}

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

    #[test]
    fn test_buffer_creation() {
        let buf = StreamBuffer::new(100, 256.0);
        assert_eq!(buf.capacity(), 100);
        assert_eq!(buf.sfreq(), 256.0);
        assert!(buf.is_empty());
    }

    #[test]
    fn test_buffer_for_duration() {
        let buf = StreamBuffer::for_duration(2.0, 256.0);
        assert_eq!(buf.capacity(), 512);
    }

    #[test]
    fn test_push_and_get() {
        let mut buf = StreamBuffer::new(10, 100.0);

        for i in 0..5 {
            buf.push_sample(vec![i as f64], i as f64 * 0.01);
        }

        assert_eq!(buf.len(), 5);

        let (data, times) = buf.get_all();
        assert_eq!(data.len(), 5);
        assert_eq!(times.len(), 5);
        assert_eq!(data[0][0], 0.0);
        assert_eq!(data[4][0], 4.0);
    }

    #[test]
    fn test_buffer_overflow() {
        let mut buf = StreamBuffer::new(5, 100.0);

        for i in 0..10 {
            buf.push_sample(vec![i as f64], i as f64 * 0.01);
        }

        assert_eq!(buf.len(), 5);
        assert_eq!(buf.dropped(), 5);
        assert_eq!(buf.total_received(), 10);

        let (data, _) = buf.get_all();
        assert_eq!(data[0][0], 5.0); // First 5 samples were dropped
    }

    #[test]
    fn test_get_recent() {
        let mut buf = StreamBuffer::new(100, 10.0);

        // Push 50 samples (5 seconds at 10 Hz)
        for i in 0..50 {
            buf.push_sample(vec![i as f64], i as f64 * 0.1);
        }

        // Get last 2 seconds (20 samples)
        let (data, _) = buf.get_recent(2.0);
        assert_eq!(data.len(), 20);
        assert_eq!(data[0][0], 30.0); // Start from sample 30
    }
}