foxhunt/crates/ml/src/batch_processing.rs

#![allow(unsafe_code)] // Intentional unsafe for SIMD batch processing performance

//! Batch Processing Optimizations for Ultra-Low Latency ML Inference
//!
//! Implements advanced batch processing techniques to achieve sub-100μs inference
//! targets for HFT applications. Features SIMD operations, memory pooling,
//! and cache-friendly data layouts.

// For canonical types

use std::collections::VecDeque;
use std::f64::consts::FRAC_2_SQRT_PI;
use std::fmt;

use ndarray::{Array1, Array2, Axis};
use serde::{Deserialize, Serialize};

use crate::{MLError, PRECISION_FACTOR};

/// Activation function types
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ActivationFunction {
    ReLU,
    LeakyReLU { alpha: f64 },
    Sigmoid,
    Tanh,
    Gelu,
}

impl fmt::Display for ActivationFunction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ActivationFunction::ReLU => write!(f, "ReLU"),
            ActivationFunction::LeakyReLU { alpha } => write!(f, "LeakyReLU(α={})", alpha),
            ActivationFunction::Sigmoid => write!(f, "Sigmoid"),
            ActivationFunction::Tanh => write!(f, "Tanh"),
            ActivationFunction::Gelu => write!(f, "GELU"),
        }
    }
}

/// Reduction operations for tensor processing
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ReductionOp {
    Sum,
    Mean,
    Max,
    Min,
}

/// Element-wise operations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ElementWiseOp {
    Add,
    Multiply,
    Subtract,
    Divide,
}

/// Configuration for batch processing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchProcessingConfig {
    pub max_batch_size: usize,
    pub use_simd: bool,
    pub memory_alignment: usize,
}

impl Default for BatchProcessingConfig {
    fn default() -> Self {
        Self {
            max_batch_size: 1024,
            use_simd: true,
            memory_alignment: 32,
        }
    }
}

/// SIMD capabilities detection
#[derive(Debug, Clone)]
pub struct SIMDCapabilities {
    pub vector_width: usize,
    pub has_avx512: bool,
    pub has_avx2: bool,
    pub has_sse4: bool,
}

impl Default for SIMDCapabilities {
    fn default() -> Self {
        Self {
            vector_width: 8, // Default to AVX2
            has_avx512: false,
            has_avx2: true,
            has_sse4: true,
        }
    }
}

/// Memory pool configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryPoolConfig {
    pub initial_capacity: usize,
    pub max_pools: usize,
    pub alignment: usize,
}

impl Default for MemoryPoolConfig {
    fn default() -> Self {
        Self {
            initial_capacity: 1024 * 1024, // 1MB
            max_pools: 16,
            alignment: 32,
        }
    }
}

/// Memory pool statistics
#[derive(Debug, Default, Clone)]
pub struct MemoryPoolStats {
    pub total_allocations: u64,
    pub total_deallocations: u64,
    pub peak_memory_usage: usize,
    pub current_memory_usage: usize,
}

/// Aligned buffer for efficient SIMD operations
#[derive(Debug)]
pub struct AlignedBuffer {
    data: Vec<f64>,
    capacity: usize,
    alignment: usize,
    len: usize,
}

impl AlignedBuffer {
    pub fn new(capacity: usize, alignment: usize) -> Result<Self, MLError> {
        if alignment == 0 || !alignment.is_power_of_two() {
            return Err(MLError::ConfigError("Alignment must be a power of two".to_owned()));
        }

        let mut data = Vec::with_capacity(capacity + alignment);
        data.resize(capacity, 0.0);

        Ok(Self {
            data,
            capacity,
            alignment,
            len: 0,
        })
    }

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

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn set_len(&mut self, len: usize) {
        if len <= self.capacity {
            self.len = len;
        }
    }

    // SAFETY: Unsafe slice access from aligned buffer
    // - Invariant 1: self.len always <= self.data.len() (enforced in with_capacity)
    // - Invariant 2: Data initialized up to self.len before calling as_slice
    // - Invariant 3: Alignment requirements preserved from allocation
    // - Verified: Constructor ensures capacity >= requested size
    // - Risk: MEDIUM - Caller must ensure data initialized before read
    pub unsafe fn as_slice(&self) -> &[f64] {
        &self.data[..self.len]
    }

    // SAFETY: Unsafe mutable slice access from aligned buffer
    // - Invariant 1: self.len <= self.data.len() (same as above)
    // - Invariant 2: Exclusive access guaranteed by &mut self
    // - Invariant 3: Slice lifetime tied to buffer, no aliasing
    // - Verified: Mutable reference ensures no concurrent reads
    // - Risk: MEDIUM - Caller must maintain slice bounds during use
    pub unsafe fn as_mut_slice(&mut self) -> &mut [f64] {
        &mut self.data[..self.len]
    }
}

/// Memory pool for efficient buffer reuse
#[derive(Debug)]
pub struct MemoryPool {
    config: MemoryPoolConfig,
    available_buffers: VecDeque<AlignedBuffer>,
    stats: MemoryPoolStats,
}

impl MemoryPool {
    pub fn new(config: MemoryPoolConfig) -> Result<Self, MLError> {
        Ok(Self {
            config,
            available_buffers: VecDeque::new(),
            stats: MemoryPoolStats::default(),
        })
    }

    pub fn get_buffer(&mut self, size: usize) -> Result<AlignedBuffer, MLError> {
        self.stats.total_allocations += 1;

        if let Some(mut buffer) = self.available_buffers.pop_front() {
            if buffer.capacity() >= size {
                buffer.set_len(size);
                return Ok(buffer);
            }
        }

        // Create new buffer
        AlignedBuffer::new(
            size.max(self.config.initial_capacity),
            self.config.alignment,
        )
    }

    pub fn return_buffer(&mut self, buffer: AlignedBuffer) {
        self.stats.total_deallocations += 1;
        if self.available_buffers.len() < self.config.max_pools {
            self.available_buffers.push_back(buffer);
        }
    }

    pub fn get_stats(&self) -> &MemoryPoolStats {
        &self.stats
    }
}

/// Auto-tuner for optimal batch sizes
#[derive(Debug)]
pub struct BatchSizeAutoTuner {
    current_batch_size: usize,
    min_batch_size: usize,
    max_batch_size: usize,
    recent_latencies: VecDeque<u64>,
    window_size: usize,
}

impl BatchSizeAutoTuner {
    pub fn new(initial_batch_size: usize) -> Self {
        Self {
            current_batch_size: initial_batch_size,
            min_batch_size: 1,
            max_batch_size: 2048,
            recent_latencies: VecDeque::new(),
            window_size: 10,
        }
    }

    pub fn update_performance(&mut self, latency_ns: u64) -> usize {
        self.recent_latencies.push_back(latency_ns);
        if self.recent_latencies.len() > self.window_size {
            self.recent_latencies.pop_front();
        }

        // Simple auto-tuning logic
        if self.recent_latencies.len() >= self.window_size {
            let avg_latency = self.recent_latencies.iter().sum::<u64>() / self.window_size as u64;

            if avg_latency > 100_000 {
                // > 100μs target
                self.current_batch_size =
                    (self.current_batch_size * 9 / 10).max(self.min_batch_size);
            } else if avg_latency < 50_000 {
                // < 50μs, can increase
                self.current_batch_size =
                    (self.current_batch_size * 11 / 10).min(self.max_batch_size);
            } else {
                // Latency within acceptable range, keep current batch size
            }
        }

        self.current_batch_size
    }
}

/// Main batch processor with optimizations
#[derive(Debug)]
pub struct BatchProcessor {
    config: BatchProcessingConfig,
    pub simd_capabilities: SIMDCapabilities,
    memory_pool: MemoryPool,
    auto_tuner: BatchSizeAutoTuner,
}

impl BatchProcessor {
    pub fn new(config: BatchProcessingConfig) -> Result<Self, MLError> {
        let memory_pool = MemoryPool::new(MemoryPoolConfig::default())?;
        let auto_tuner = BatchSizeAutoTuner::new(config.max_batch_size / 4);

        Ok(Self {
            config,
            simd_capabilities: SIMDCapabilities::default(),
            memory_pool,
            auto_tuner,
        })
    }

    /// Standard matrix multiplication
    pub fn standard_matrix_multiply(
        a: &Array2<i64>,
        b: &Array2<i64>,
    ) -> Result<Array2<i64>, MLError> {
        if a.ncols() != b.nrows() {
            return Err(MLError::DimensionMismatch {
                expected: a.ncols(),
                actual: b.nrows(),
            });
        }

        let result = a.dot(b);
        Ok(result)
    }

    /// Standard element-wise operations
    pub fn standard_element_wise_operation(
        op: &ElementWiseOp,
        inputs: &[Array1<i64>],
    ) -> Result<Array1<i64>, MLError> {
        if inputs.is_empty() {
            return Err(MLError::InvalidInput(
                "No input arrays provided".to_owned(),
            ));
        }

        // Pre-allocate result array to avoid clone
        let len = inputs[0].len();
        let mut result = Array1::zeros(len);
        result.assign(&inputs[0]);

        for input in &inputs[1..] {
            if input.len() != result.len() {
                return Err(MLError::DimensionMismatch {
                    expected: result.len(),
                    actual: input.len(),
                });
            }

            match op {
                ElementWiseOp::Add => {
                    for i in 0..result.len() {
                        if let Some(val) = input.get(i) {
                            result[i] += val;
                        }
                    }
                },
                ElementWiseOp::Multiply => {
                    for i in 0..result.len() {
                        if let Some(val) = input.get(i) {
                            result[i] = (result[i] * val) / PRECISION_FACTOR as i64;
                        }
                    }
                },
                ElementWiseOp::Subtract => {
                    for i in 0..result.len() {
                        if let Some(val) = input.get(i) {
                            result[i] -= val;
                        }
                    }
                },
                ElementWiseOp::Divide => {
                    for i in 0..result.len() {
                        if let Some(&val) = input.get(i) {
                            if val != 0 {
                                result[i] = (result[i] * PRECISION_FACTOR as i64) / val;
                            }
                        }
                    }
                },
            }
        }

        Ok(result)
    }

    /// Standard activation functions
    pub fn standard_activation(
        input: &Array1<i64>,
        activation: &ActivationFunction,
    ) -> Result<Array1<i64>, MLError> {
        let mut result = Array1::zeros(input.len());

        for i in 0..input.len() {
            let val = input.get(i).copied().ok_or_else(|| {
                MLError::InvalidInput(format!(
                    "Index {} out of bounds (length {})",
                    i,
                    input.len()
                ))
            })?;
            let x = val as f64 / PRECISION_FACTOR as f64;
            let activated = match activation {
                ActivationFunction::ReLU => x.max(0.0),
                ActivationFunction::LeakyReLU { alpha } => {
                    if x > 0.0 {
                        x
                    } else {
                        alpha * x
                    }
                },
                ActivationFunction::Sigmoid => 1.0 / (1.0 + (-x).exp()),
                ActivationFunction::Tanh => x.tanh(),
                ActivationFunction::Gelu => {
                    0.5 * x * (1.0 + (x * FRAC_2_SQRT_PI * 0.7978845608).tanh())
                },
            };
            result[i] = (activated * PRECISION_FACTOR as f64) as i64;
        }

        Ok(result)
    }

    /// Standard reduction operations
    pub fn reduction_operation(
        input: &Array2<i64>,
        op: &ReductionOp,
        axis: Option<usize>,
    ) -> Result<Array1<i64>, MLError> {
        match op {
            ReductionOp::Sum => {
                if let Some(ax) = axis {
                    if ax >= input.ndim() {
                        return Err(MLError::DimensionMismatch {
                            expected: input.ndim(),
                            actual: ax,
                        });
                    }
                    Ok(input.sum_axis(Axis(ax)))
                } else {
                    let total_sum = input.sum();
                    Ok(Array1::from_elem(1, total_sum))
                }
            },
            ReductionOp::Mean => {
                if let Some(ax) = axis {
                    if ax >= input.ndim() {
                        return Err(MLError::DimensionMismatch {
                            expected: input.ndim(),
                            actual: ax,
                        });
                    }
                    let sum = input.sum_axis(Axis(ax));
                    let count = input.len_of(Axis(ax)) as i64;
                    Ok(sum.mapv(|x| x / count))
                } else {
                    let mean = input.sum() / input.len() as i64;
                    Ok(Array1::from_elem(1, mean))
                }
            },
            ReductionOp::Max => {
                if let Some(ax) = axis {
                    if ax >= input.ndim() {
                        return Err(MLError::DimensionMismatch {
                            expected: input.ndim(),
                            actual: ax,
                        });
                    }
                    Ok(input.map_axis(Axis(ax), |lane| *lane.iter().max().unwrap_or(&0)))
                } else {
                    let max_val = input.iter().max().copied().unwrap_or(0);
                    Ok(Array1::from_elem(1, max_val))
                }
            },
            ReductionOp::Min => {
                if let Some(ax) = axis {
                    if ax >= input.ndim() {
                        return Err(MLError::DimensionMismatch {
                            expected: input.ndim(),
                            actual: ax,
                        });
                    }
                    Ok(input.map_axis(Axis(ax), |lane| *lane.iter().min().unwrap_or(&0)))
                } else {
                    let min_val = input.iter().min().copied().unwrap_or(0);
                    Ok(Array1::from_elem(1, min_val))
                }
            },
        }
    }
}

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

    #[test]
    fn test_batch_processor_creation() {
        let config = BatchProcessingConfig::default();
        let processor = BatchProcessor::new(config).unwrap();
        assert!(processor.simd_capabilities.vector_width > 0);
    }

    #[test]
    fn test_aligned_buffer() {
        let mut buffer = AlignedBuffer::new(1024, 32).unwrap();
        assert_eq!(buffer.capacity(), 1024);
        buffer.set_len(512);
        assert_eq!(buffer.len(), 512);
    }

    #[test]
    fn test_aligned_buffer_invalid_alignment() {
        // Non-power-of-two alignment should fail
        let result = AlignedBuffer::new(1024, 31);
        assert!(result.is_err());

        // Zero alignment should fail
        let result = AlignedBuffer::new(1024, 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_memory_pool() {
        let config = MemoryPoolConfig::default();
        let mut pool = MemoryPool::new(config).unwrap();

        // Get buffer
        let buffer = pool.get_buffer(256).unwrap();
        assert!(buffer.capacity() >= 256);

        // Return buffer
        pool.return_buffer(buffer);

        let stats = pool.get_stats();
        assert_eq!(stats.total_allocations, 1);
        assert_eq!(stats.total_deallocations, 1);
    }

    #[test]
    fn test_memory_pool_reuse() {
        let mut pool = MemoryPool::new(MemoryPoolConfig::default()).unwrap();

        // Get and return buffer
        let buffer1 = pool.get_buffer(512).unwrap();
        pool.return_buffer(buffer1);

        // Get another buffer - should reuse
        let _buffer2 = pool.get_buffer(512).unwrap();

        let stats = pool.get_stats();
        assert_eq!(stats.total_allocations, 2);
    }

    #[test]
    fn test_matrix_multiply() {
        let a = Array2::from_shape_vec((2, 3), vec![1, 2, 3, 4, 5, 6]).unwrap();
        let b = Array2::from_shape_vec((3, 2), vec![7, 8, 9, 10, 11, 12]).unwrap();

        let result = BatchProcessor::standard_matrix_multiply(&a, &b).unwrap();
        assert_eq!(result.dim(), (2, 2));
        // Verify correctness: [1,2,3] * [7,9,11; 8,10,12] = [58, 64]
        assert_eq!(result[[0, 0]], 58);
        assert_eq!(result[[0, 1]], 64);
    }

    #[test]
    fn test_matrix_multiply_dimension_mismatch() {
        let a = Array2::from_shape_vec((2, 3), vec![1, 2, 3, 4, 5, 6]).unwrap();
        let b = Array2::from_shape_vec((2, 2), vec![7, 8, 9, 10]).unwrap();

        let result = BatchProcessor::standard_matrix_multiply(&a, &b);
        assert!(result.is_err());
    }

    #[test]
    fn test_element_wise_add() {
        let a = Array1::from_vec(vec![100, 200, 300]);
        let b = Array1::from_vec(vec![50, 100, 150]);

        let result =
            BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Add, &[a, b]).unwrap();

        assert_eq!(result, Array1::from_vec(vec![150, 300, 450]));
    }

    #[test]
    fn test_element_wise_multiply() {
        let a = Array1::from_vec(vec![100_000_000, 200_000_000, 300_000_000]);
        let b = Array1::from_vec(vec![200_000_000, 300_000_000, 400_000_000]);

        let result =
            BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Multiply, &[a, b])
                .unwrap();

        // Result: 2.0, 6.0, 12.0 in fixed point
        assert_eq!(result.get(0).copied().unwrap(), 200_000_000);
        assert_eq!(result.get(1).copied().unwrap(), 600_000_000);
        assert_eq!(result.get(2).copied().unwrap(), 1_200_000_000);
    }

    #[test]
    fn test_element_wise_subtract() {
        let a = Array1::from_vec(vec![100, 200, 300]);
        let b = Array1::from_vec(vec![50, 100, 150]);

        let result =
            BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Subtract, &[a, b])
                .unwrap();

        assert_eq!(result, Array1::from_vec(vec![50, 100, 150]));
    }

    #[test]
    fn test_element_wise_divide() {
        let a = Array1::from_vec(vec![200_000_000, 600_000_000, 1_200_000_000]);
        let b = Array1::from_vec(vec![100_000_000, 200_000_000, 300_000_000]);

        let result =
            BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Divide, &[a, b])
                .unwrap();

        assert_eq!(result.get(0).copied().unwrap(), 200_000_000); // 2.0
        assert_eq!(result.get(1).copied().unwrap(), 300_000_000); // 3.0
        assert_eq!(result.get(2).copied().unwrap(), 400_000_000); // 4.0
    }

    #[test]
    fn test_element_wise_divide_by_zero() {
        let a = Array1::from_vec(vec![100_000_000, 200_000_000]);
        let b = Array1::from_vec(vec![0, 100_000_000]);

        let result =
            BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Divide, &[a, b])
                .unwrap();

        // Division by zero protection
        assert_eq!(result.get(0).copied().unwrap(), 100_000_000);
        assert_eq!(result.get(1).copied().unwrap(), 200_000_000);
    }

    #[test]
    fn test_element_wise_empty_inputs() {
        let result = BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Add, &[]);
        assert!(result.is_err());
    }

    #[test]
    fn test_element_wise_dimension_mismatch() {
        let a = Array1::from_vec(vec![100, 200, 300]);
        let b = Array1::from_vec(vec![50, 100]); // Different size

        let result = BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Add, &[a, b]);
        assert!(result.is_err());
    }

    #[test]
    fn test_batch_size_auto_tuner() {
        let mut tuner = BatchSizeAutoTuner::new(32);

        // Simulate high latency - should decrease batch size
        for _ in 0..11 {
            tuner.update_performance(200_000); // 200μs
        }
        let new_size = tuner.update_performance(200_000);
        assert!(new_size < 32);

        // Simulate low latency - should increase batch size
        // Need extra iterations to flush the sliding window and allow size to increase
        for _ in 0..24 {
            tuner.update_performance(30_000); // 30μs
        }
        let final_size = tuner.update_performance(30_000);
        assert!(final_size > 32);
    }

    #[test]
    fn test_batch_size_auto_tuner_bounds() {
        let mut tuner = BatchSizeAutoTuner::new(1);

        // Try to decrease below minimum
        for _ in 0..15 {
            tuner.update_performance(200_000);
        }
        let size = tuner.update_performance(200_000);
        assert!(size >= tuner.min_batch_size);
    }

    #[test]
    fn test_activation_function_display() {
        assert_eq!(format!("{}", ActivationFunction::ReLU), "ReLU");
        assert_eq!(format!("{}", ActivationFunction::Sigmoid), "Sigmoid");
        assert_eq!(format!("{}", ActivationFunction::Tanh), "Tanh");
        assert_eq!(format!("{}", ActivationFunction::Gelu), "GELU");
        assert_eq!(
            format!("{}", ActivationFunction::LeakyReLU { alpha: 0.01 }),
            "LeakyReLU(α=0.01)"
        );
    }

    #[test]
    fn test_batch_processing_config_default() {
        let config = BatchProcessingConfig::default();
        assert_eq!(config.max_batch_size, 1024);
        assert!(config.use_simd);
        assert_eq!(config.memory_alignment, 32);
    }

    #[test]
    fn test_simd_capabilities_default() {
        let caps = SIMDCapabilities::default();
        assert_eq!(caps.vector_width, 8);
        assert!(caps.has_avx2);
        assert!(caps.has_sse4);
    }
}