#![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 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 std::fmt::Display for ActivationFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::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 {
                reason: "Alignment must be a power of two".to_string(),
            });
        }

        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;
        }
    }

    pub unsafe fn as_slice(&self) -> &[f64] {
        &self.data[..self.len]
    }

    pub unsafe fn as_mut_slice(&mut self) -> &mut [f64] {
        &mut self.data[..self.len]
    }
}

/// Memory pool for efficient buffer reuse
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.clone()
    }
}

/// Auto-tuner for optimal batch sizes
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);
            }
        }

        self.current_batch_size
    }
}

/// Main batch processor with optimizations
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_string(),
            ));
        }

        let mut result = inputs[0].clone();

        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() {
                        result[i] += input[i];
                    }
                }
                ElementWiseOp::Multiply => {
                    for i in 0..result.len() {
                        result[i] = (result[i] * input[i]) / PRECISION_FACTOR as i64;
                    }
                }
                ElementWiseOp::Subtract => {
                    for i in 0..result.len() {
                        result[i] -= input[i];
                    }
                }
                ElementWiseOp::Divide => {
                    for i in 0..result.len() {
                        if input[i] != 0 {
                            result[i] = (result[i] * PRECISION_FACTOR as i64) / input[i];
                        }
                    }
                }
            }
        }

        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 x = input[i] 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 * std::f64::consts::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::*;
    use ndarray::array;
    // use crate::safe_operations; // DISABLED - module not found

    #[test]
    fn test_batch_processor_creation() {
        let config = BatchProcessingConfig::default();
        let processor = BatchProcessor::new(config);

        assert!(processor.is_ok());
        let processor = processor?;
        assert!(processor.simd_capabilities.vector_width > 0);
    }

    #[test]
    fn test_aligned_buffer() {
        let buffer = AlignedBuffer::new(1024, 32);

        assert!(buffer.is_ok());
        let mut buffer = buffer?;
        assert_eq!(buffer.capacity(), 1024);

        buffer.set_len(512);
        unsafe {
            let slice = buffer.as_slice();
            assert_eq!(slice.len(), 512);
        }
    }

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

        assert!(pool.is_ok());
        let mut pool = pool?;

        // Get buffer
        let buffer = pool.get_buffer(256);
        assert!(buffer.is_ok());
        let buffer = buffer?;
        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_matrix_multiply() {
        let a = array![[1000, 2000], [3000, 4000]]; // 0.1, 0.2; 0.3, 0.4 in fixed-point
        let b = array![[5000, 6000], [7000, 8000]]; // 0.5, 0.6; 0.7, 0.8 in fixed-point

        let result = BatchProcessor::standard_matrix_multiply(&a, &b);

        assert!(result.is_ok());
        let result = result?;

        // Expected result: [[0.1*0.5 + 0.2*0.7, 0.1*0.6 + 0.2*0.8], [0.3*0.5 + 0.4*0.7, 0.3*0.6 + 0.4*0.8]]
        // = [[0.19, 0.22], [0.43, 0.50]]
        assert_eq!(result[[0, 0]], 1900); // 0.19 in fixed-point
        assert_eq!(result[[0, 1]], 2200); // 0.22 in fixed-point
        assert_eq!(result[[1, 0]], 4300); // 0.43 in fixed-point
        assert_eq!(result[[1, 1]], 5000); // 0.50 in fixed-point
    }

    #[test]
    fn test_element_wise_operations() {
        let input1 = array![1000, 2000, 3000]; // 0.1, 0.2, 0.3 in fixed-point
        let input2 = array![4000, 5000, 6000]; // 0.4, 0.5, 0.6 in fixed-point
        let inputs = vec![input1, input2];

        // Test addition
        let result = BatchProcessor::standard_element_wise_operation(&ElementWiseOp::Add, &inputs);
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 5000); // 0.5 in fixed-point
        assert_eq!(result[1], 7000); // 0.7 in fixed-point
        assert_eq!(result[2], 9000); // 0.9 in fixed-point
    }

    #[test]
    fn test_activation_functions() {
        let input = array![1000, -2000, 3000]; // 0.1, -0.2, 0.3 in fixed-point

        // Test ReLU
        let result = BatchProcessor::standard_activation(&input, &ActivationFunction::ReLU);
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 1000); // 0.1
        assert_eq!(result[1], 0); // 0.0 (ReLU clips negative)
        assert_eq!(result[2], 3000); // 0.3

        // Test LeakyReLU
        let alpha = 0.1;
        let result =
            BatchProcessor::standard_activation(&input, &ActivationFunction::LeakyReLU { alpha });
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 1000); // 0.1
        assert_eq!(result[1], -200); // -0.02 (alpha * -0.2)
        assert_eq!(result[2], 3000); // 0.3
    }

    #[test]
    fn test_reduction_operations() {
        let input = array![[1000, 2000], [3000, 4000]]; // [[0.1, 0.2], [0.3, 0.4]]

        // Test sum along axis 0
        let result = BatchProcessor::reduction_operation(&input, &ReductionOp::Sum, Some(0));
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 4000); // 0.1 + 0.3 = 0.4
        assert_eq!(result[1], 6000); // 0.2 + 0.4 = 0.6

        // Test sum along axis 1
        let result = BatchProcessor::reduction_operation(&input, &ReductionOp::Sum, Some(1));
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 3000); // 0.1 + 0.2 = 0.3
        assert_eq!(result[1], 7000); // 0.3 + 0.4 = 0.7

        // Test max along axis 0
        let result = BatchProcessor::reduction_operation(&input, &ReductionOp::Max, Some(0));
        assert!(result.is_ok());
        let result = result?;
        assert_eq!(result[0], 3000); // max(0.1, 0.3) = 0.3
        assert_eq!(result[1], 4000); // max(0.2, 0.4) = 0.4
    }

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

        // Simulate high latency - should decrease batch size
        let new_size = tuner.update_performance(200_000); // 200μs

        // Should have decreased from initial size
        assert!(new_size <= 32);

        // Simulate low latency - should increase batch size
        for _ in 0..10 {
            tuner.update_performance(30_000); // 30μs
        }
        let final_size = tuner.update_performance(30_000);

        // Should have increased due to low latencies
        assert!(final_size > 32);
    }
}