rustytorch/crates/specialized/rtx-neuro-realtime/src/latency.rs

//! Latency monitoring for real-time pipeline stages.
//!
//! Tracks per-stage latency to identify bottlenecks and ensure
//! the pipeline meets sub-10ms requirements.

use parking_lot::RwLock;
use std::collections::VecDeque;
use std::time::{Duration, Instant};

/// Processing stage for latency tracking
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LatencyStage {
    /// LSL data receive
    LslReceive,
    /// GPU filter application
    GpuFilter,
    /// GPU beamformer computation
    GpuBeamformer,
    /// Source estimate computation
    SourceEstimate,
    /// Total end-to-end
    Total,
}

impl std::fmt::Display for LatencyStage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::LslReceive => write!(f, "LSL Receive"),
            Self::GpuFilter => write!(f, "GPU Filter"),
            Self::GpuBeamformer => write!(f, "GPU Beamformer"),
            Self::SourceEstimate => write!(f, "Source Estimate"),
            Self::Total => write!(f, "Total"),
        }
    }
}

/// Latency statistics for a stage
#[derive(Debug, Clone)]
pub struct LatencyStats {
    /// Stage name
    pub stage: LatencyStage,
    /// Mean latency in microseconds
    pub mean_us: f64,
    /// Standard deviation in microseconds
    pub std_us: f64,
    /// Minimum latency in microseconds
    pub min_us: f64,
    /// Maximum latency in microseconds
    pub max_us: f64,
    /// 50th percentile (median) in microseconds
    pub p50_us: f64,
    /// 95th percentile in microseconds
    pub p95_us: f64,
    /// 99th percentile in microseconds
    pub p99_us: f64,
    /// Number of samples
    pub count: usize,
}

impl LatencyStats {
    /// Mean latency in milliseconds
    pub fn mean_ms(&self) -> f64 {
        self.mean_us / 1000.0
    }

    /// 95th percentile in milliseconds
    pub fn p95_ms(&self) -> f64 {
        self.p95_us / 1000.0
    }

    /// 99th percentile in milliseconds
    pub fn p99_ms(&self) -> f64 {
        self.p99_us / 1000.0
    }

    /// Check if latency meets target
    pub fn meets_target(&self, target_ms: f64) -> bool {
        self.p99_ms() <= target_ms
    }
}

/// Per-stage latency tracker
#[derive(Debug)]
struct StageTracker {
    /// Recent latency samples in microseconds
    samples: VecDeque<f64>,
    /// Maximum samples to keep
    capacity: usize,
    /// Current timing start
    start: Option<Instant>,
}

impl StageTracker {
    fn new(capacity: usize) -> Self {
        Self {
            samples: VecDeque::with_capacity(capacity),
            capacity,
            start: None,
        }
    }

    fn start(&mut self) {
        self.start = Some(Instant::now());
    }

    fn stop(&mut self) {
        if let Some(start) = self.start.take() {
            let duration_us = start.elapsed().as_secs_f64() * 1_000_000.0;
            if self.samples.len() >= self.capacity {
                self.samples.pop_front();
            }
            self.samples.push_back(duration_us);
        }
    }

    fn record(&mut self, duration: Duration) {
        let duration_us = duration.as_secs_f64() * 1_000_000.0;
        if self.samples.len() >= self.capacity {
            self.samples.pop_front();
        }
        self.samples.push_back(duration_us);
    }

    fn stats(&self, stage: LatencyStage) -> LatencyStats {
        if self.samples.is_empty() {
            return LatencyStats {
                stage,
                mean_us: 0.0,
                std_us: 0.0,
                min_us: 0.0,
                max_us: 0.0,
                p50_us: 0.0,
                p95_us: 0.0,
                p99_us: 0.0,
                count: 0,
            };
        }

        let n = self.samples.len();
        let sum: f64 = self.samples.iter().sum();
        let mean = sum / n as f64;

        let variance: f64 = self.samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n as f64;
        let std = variance.sqrt();

        let min = self.samples.iter().copied().fold(f64::INFINITY, f64::min);
        let max = self
            .samples
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max);

        // Sort for percentiles
        let mut sorted: Vec<f64> = self.samples.iter().copied().collect();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let percentile = |p: f64| -> f64 {
            let idx = ((p / 100.0) * (n - 1) as f64).round() as usize;
            sorted[idx.min(n - 1)]
        };

        LatencyStats {
            stage,
            mean_us: mean,
            std_us: std,
            min_us: min,
            max_us: max,
            p50_us: percentile(50.0),
            p95_us: percentile(95.0),
            p99_us: percentile(99.0),
            count: n,
        }
    }

    fn clear(&mut self) {
        self.samples.clear();
        self.start = None;
    }
}

/// Latency monitor for the real-time pipeline
#[derive(Debug)]
pub struct LatencyMonitor {
    /// Per-stage trackers
    trackers: RwLock<[StageTracker; 5]>,
    /// Target latencies in milliseconds
    targets: [f64; 5],
    /// Monitoring window size
    window_size: usize,
}

impl LatencyMonitor {
    /// Create a new latency monitor
    ///
    /// # Arguments
    /// * `window_size` - Number of samples to keep for statistics
    pub fn new(window_size: usize) -> Self {
        Self {
            trackers: RwLock::new([
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
            ]),
            targets: [1.0, 0.5, 2.0, 1.0, 5.0], // Default targets in ms
            window_size,
        }
    }

    /// Create with custom target latencies
    pub fn with_targets(window_size: usize, targets: [f64; 5]) -> Self {
        Self {
            trackers: RwLock::new([
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
                StageTracker::new(window_size),
            ]),
            targets,
            window_size,
        }
    }

    fn stage_index(stage: LatencyStage) -> usize {
        match stage {
            LatencyStage::LslReceive => 0,
            LatencyStage::GpuFilter => 1,
            LatencyStage::GpuBeamformer => 2,
            LatencyStage::SourceEstimate => 3,
            LatencyStage::Total => 4,
        }
    }

    /// Start timing a stage
    pub fn start(&self, stage: LatencyStage) {
        let idx = Self::stage_index(stage);
        self.trackers.write()[idx].start();
    }

    /// Stop timing a stage
    pub fn stop(&self, stage: LatencyStage) {
        let idx = Self::stage_index(stage);
        self.trackers.write()[idx].stop();
    }

    /// Record a duration directly
    pub fn record(&self, stage: LatencyStage, duration: Duration) {
        let idx = Self::stage_index(stage);
        self.trackers.write()[idx].record(duration);
    }

    /// Get statistics for a stage
    pub fn stats(&self, stage: LatencyStage) -> LatencyStats {
        let idx = Self::stage_index(stage);
        self.trackers.read()[idx].stats(stage)
    }

    /// Get statistics for all stages
    pub fn all_stats(&self) -> Vec<LatencyStats> {
        let trackers = self.trackers.read();
        vec![
            trackers[0].stats(LatencyStage::LslReceive),
            trackers[1].stats(LatencyStage::GpuFilter),
            trackers[2].stats(LatencyStage::GpuBeamformer),
            trackers[3].stats(LatencyStage::SourceEstimate),
            trackers[4].stats(LatencyStage::Total),
        ]
    }

    /// Check if a stage meets its target
    pub fn meets_target(&self, stage: LatencyStage) -> bool {
        let idx = Self::stage_index(stage);
        let stats = self.trackers.read()[idx].stats(stage);
        stats.p99_ms() <= self.targets[idx]
    }

    /// Check if all stages meet their targets
    pub fn all_meet_targets(&self) -> bool {
        (0..5).all(|i| {
            let stage = match i {
                0 => LatencyStage::LslReceive,
                1 => LatencyStage::GpuFilter,
                2 => LatencyStage::GpuBeamformer,
                3 => LatencyStage::SourceEstimate,
                _ => LatencyStage::Total,
            };
            self.meets_target(stage)
        })
    }

    /// Get target latency for a stage
    pub fn target(&self, stage: LatencyStage) -> f64 {
        self.targets[Self::stage_index(stage)]
    }

    /// Clear all statistics
    pub fn clear(&self) {
        for tracker in self.trackers.write().iter_mut() {
            tracker.clear();
        }
    }

    /// Generate a summary report
    pub fn report(&self) -> String {
        let stats = self.all_stats();
        let mut lines = vec!["Latency Report".to_string(), "=".repeat(60)];

        for (i, s) in stats.iter().enumerate() {
            let target = self.targets[i];
            let meets = s.meets_target(target);
            let status = if meets { "OK" } else { "VIOLATION" };

            lines.push(format!(
                "{:20} | mean: {:7.2} us | p99: {:7.2} us | target: {:5.1} ms | {}",
                s.stage.to_string(),
                s.mean_us,
                s.p99_us,
                target,
                status
            ));
        }

        lines.join("\n")
    }
}

impl Default for LatencyMonitor {
    fn default() -> Self {
        Self::new(1000) // Keep last 1000 samples
    }
}

/// RAII timing guard for automatic stage timing
pub struct TimingGuard<'a> {
    monitor: &'a LatencyMonitor,
    stage: LatencyStage,
}

impl<'a> TimingGuard<'a> {
    /// Create a new timing guard (starts timing immediately)
    pub fn new(monitor: &'a LatencyMonitor, stage: LatencyStage) -> Self {
        monitor.start(stage);
        Self { monitor, stage }
    }
}

impl Drop for TimingGuard<'_> {
    fn drop(&mut self) {
        self.monitor.stop(self.stage);
    }
}

/// Helper macro for timing a code block
#[macro_export]
macro_rules! time_stage {
    ($monitor:expr, $stage:expr, $code:block) => {{
        let _guard = $crate::latency::TimingGuard::new($monitor, $stage);
        $code
    }};
}

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

    #[test]
    fn test_stage_tracker() {
        let mut tracker = StageTracker::new(100);

        tracker.record(Duration::from_micros(100));
        tracker.record(Duration::from_micros(200));
        tracker.record(Duration::from_micros(300));

        let stats = tracker.stats(LatencyStage::Total);
        assert_eq!(stats.count, 3);
        assert!((stats.mean_us - 200.0).abs() < 1.0);
        assert!((stats.min_us - 100.0).abs() < 1.0);
        assert!((stats.max_us - 300.0).abs() < 1.0);
    }

    #[test]
    fn test_latency_monitor() {
        let monitor = LatencyMonitor::new(100);

        for _ in 0..10 {
            monitor.start(LatencyStage::GpuFilter);
            thread::sleep(Duration::from_micros(100));
            monitor.stop(LatencyStage::GpuFilter);
        }

        let stats = monitor.stats(LatencyStage::GpuFilter);
        assert_eq!(stats.count, 10);
        assert!(stats.mean_us > 50.0); // Should be at least 100us
    }

    #[test]
    fn test_timing_guard() {
        let monitor = LatencyMonitor::new(100);

        {
            let _guard = TimingGuard::new(&monitor, LatencyStage::Total);
            thread::sleep(Duration::from_micros(500));
        }

        let stats = monitor.stats(LatencyStage::Total);
        assert_eq!(stats.count, 1);
        assert!(stats.mean_us > 400.0);
    }

    #[test]
    fn test_meets_target() {
        let monitor = LatencyMonitor::with_targets(100, [1.0, 0.5, 2.0, 1.0, 5.0]);

        // Record some fast operations
        for _ in 0..10 {
            monitor.record(LatencyStage::GpuFilter, Duration::from_micros(100));
        }

        // 100us = 0.1ms, which is below target of 0.5ms
        assert!(monitor.meets_target(LatencyStage::GpuFilter));
    }

    #[test]
    fn test_report() {
        let monitor = LatencyMonitor::new(100);

        monitor.record(LatencyStage::LslReceive, Duration::from_micros(500));
        monitor.record(LatencyStage::GpuFilter, Duration::from_micros(200));
        monitor.record(LatencyStage::GpuBeamformer, Duration::from_micros(1000));
        monitor.record(LatencyStage::SourceEstimate, Duration::from_micros(300));
        monitor.record(LatencyStage::Total, Duration::from_micros(2000));

        let report = monitor.report();
        assert!(report.contains("Latency Report"));
        assert!(report.contains("GPU Filter"));
    }
}