scirs2-io 0.4.2

//! GPU-accelerated I/O operations with comprehensive backend support
//!
//! This module provides a complete GPU acceleration framework for I/O operations
//! including backend management, compression, memory management, and performance
//! optimization across CUDA, Metal, and OpenCL backends.

pub mod backend_management;
pub mod compression;
pub mod memory_management;

// Re-export key types for easy access
pub use backend_management::{
    BackendCapabilities, BackendPerformanceProfile, GpuIoProcessor, GpuWorkloadType,
};

pub use compression::{CompressionStats, GpuCompressionProcessor};

pub use memory_management::{
    AdvancedGpuMemoryPool, AllocationStats, BufferMetadata, FragmentationManager, GlobalPoolStats,
    GpuMemoryPoolManager, MemoryType, PoolConfig, PoolStats, PooledBuffer,
};

use crate::error::Result;
use scirs2_core::gpu::{GpuBackend, GpuDataType, GpuDevice};
use scirs2_core::ndarray::{Array1, ArrayView1};

/// Unified GPU I/O interface combining all GPU acceleration capabilities
#[derive(Debug)]
pub struct UnifiedGpuProcessor {
    io_processor: GpuIoProcessor,
    compression_processor: GpuCompressionProcessor,
    memory_manager: GpuMemoryPoolManager,
}

impl UnifiedGpuProcessor {
    /// Create a new unified GPU processor with optimal configuration
    pub fn new() -> Result<Self> {
        let io_processor = GpuIoProcessor::new()?;
        let compression_processor = GpuCompressionProcessor::new()?;
        let memory_manager = GpuMemoryPoolManager::new(io_processor.device.clone())?;

        Ok(Self {
            io_processor,
            compression_processor,
            memory_manager,
        })
    }

    /// Create with specific backend
    pub fn with_backend(backend: GpuBackend) -> Result<Self> {
        let io_processor = GpuIoProcessor::with_backend(backend)?;
        let compression_processor = GpuCompressionProcessor::new()?;
        let memory_manager = GpuMemoryPoolManager::new(io_processor.device.clone())?;

        Ok(Self {
            io_processor,
            compression_processor,
            memory_manager,
        })
    }

    /// Get the current GPU backend
    pub fn backend(&self) -> GpuBackend {
        self.io_processor.backend()
    }

    /// Get comprehensive GPU capabilities
    pub fn get_capabilities(&self) -> Result<GpuCapabilities> {
        let backend_caps = self.io_processor.get_backend_capabilities()?;
        let compression_stats = self.compression_processor.get_performance_stats();
        let memory_stats = self.memory_manager.get_global_stats();

        Ok(GpuCapabilities {
            backend: backend_caps.backend,
            memory_gb: backend_caps.memory_gb,
            compute_units: backend_caps.compute_units,
            supports_fp64: backend_caps.supports_fp64,
            supports_fp16: backend_caps.supports_fp16,
            compression_throughput_gbps: compression_stats.estimated_throughput_gbps,
            memory_pools: memory_stats.pool_count,
            total_pool_size: memory_stats.total_pool_size,
            performance_score: self.calculate_performance_score(&backend_caps),
        })
    }

    /// Compress data with automatic backend optimization
    pub fn compress<T: GpuDataType>(
        &self,
        data: &ArrayView1<T>,
        algorithm: crate::compression::CompressionAlgorithm,
        level: Option<u32>,
    ) -> Result<Vec<u8>> {
        self.compression_processor
            .compress_gpu(data, algorithm, level)
    }

    /// Decompress data with automatic backend optimization
    pub fn decompress<T: GpuDataType>(
        &self,
        compressed_data: &[u8],
        algorithm: crate::compression::CompressionAlgorithm,
        expected_size: usize,
    ) -> Result<Array1<T>> {
        self.compression_processor
            .decompress_gpu(compressed_data, algorithm, expected_size)
    }

    /// Allocate GPU memory buffer
    pub fn allocate_buffer(
        &mut self,
        size: usize,
        memory_type: MemoryType,
    ) -> Result<PooledBuffer> {
        self.memory_manager.allocate(size, memory_type)
    }

    /// Return buffer to memory pool
    pub fn deallocate_buffer(
        &mut self,
        buffer: PooledBuffer,
        memory_type: MemoryType,
    ) -> Result<()> {
        self.memory_manager.deallocate(buffer, memory_type)
    }

    /// Get comprehensive performance statistics
    pub fn get_performance_stats(&self) -> UnifiedGpuStats {
        let backend_caps = self
            .io_processor
            .get_backend_capabilities()
            .unwrap_or_else(|_| BackendCapabilities {
                backend: GpuBackend::Cpu,
                memory_gb: 1.0,
                max_work_group_size: 64,
                supports_fp64: false,
                supports_fp16: false,
                compute_units: 1,
                max_allocation_size: 1024 * 1024,
                local_memory_size: 64 * 1024,
            });

        let compression_stats = self.compression_processor.get_performance_stats();
        let memory_stats = self.memory_manager.get_global_stats();

        UnifiedGpuStats {
            backend: backend_caps.backend,
            compression_stats,
            memory_stats: memory_stats.clone(),
            overall_efficiency: self.calculate_efficiency_score(&memory_stats),
        }
    }

    /// Perform maintenance operations (garbage collection, compaction)
    pub fn maintenance(&mut self) -> Result<MaintenanceReport> {
        let freed_buffers = self.memory_manager.garbage_collect_all()?;

        Ok(MaintenanceReport {
            freed_buffers,
            timestamp: std::time::Instant::now(),
        })
    }

    /// Optimize processor for specific workload
    pub fn optimize_for_workload(&mut self, workload: GpuWorkloadType) -> Result<()> {
        // Adjust memory pool configurations based on workload
        match workload {
            GpuWorkloadType::MachineLearning => {
                // ML workloads benefit from larger buffers and device memory
                let pool_size = 512 * 1024 * 1024; // 512MB
                self.memory_manager
                    .create_pool(pool_size, MemoryType::Device)?;
            }
            GpuWorkloadType::ImageProcessing => {
                // Image processing needs unified memory for CPU/GPU transfers
                let pool_size = 256 * 1024 * 1024; // 256MB
                self.memory_manager
                    .create_pool(pool_size, MemoryType::Unified)?;
            }
            GpuWorkloadType::Compression => {
                // Compression benefits from pinned memory for fast transfers
                let pool_size = 128 * 1024 * 1024; // 128MB
                self.memory_manager
                    .create_pool(pool_size, MemoryType::Pinned)?;
            }
            GpuWorkloadType::GeneralCompute => {
                // Balanced approach for general compute
                let pool_size = 256 * 1024 * 1024; // 256MB
                self.memory_manager
                    .create_pool(pool_size, MemoryType::Device)?;
            }
        }

        Ok(())
    }

    // Private helper methods
    fn calculate_performance_score(&self, caps: &BackendCapabilities) -> f64 {
        let memory_score = (caps.memory_gb / 16.0).min(1.0); // Normalize to 16GB max
        let compute_score = (caps.compute_units as f64 / 64.0).min(1.0); // Normalize to 64 units
        let feature_score = if caps.supports_fp64 && caps.supports_fp16 {
            1.0
        } else {
            0.7
        };

        (memory_score + compute_score + feature_score) / 3.0
    }

    fn calculate_efficiency_score(&self, memory_stats: &GlobalPoolStats) -> f64 {
        let allocation_efficiency = memory_stats.global_allocation_stats.get_cache_hit_rate();
        let fragmentation_penalty = 1.0 - memory_stats.average_fragmentation.min(1.0);

        (allocation_efficiency + fragmentation_penalty) / 2.0
    }
}

impl Default for UnifiedGpuProcessor {
    fn default() -> Self {
        Self::new().unwrap_or_else(|_| {
            // Fallback to CPU-only processor
            let device = GpuDevice::new(GpuBackend::Cpu, 0);
            Self {
                io_processor: GpuIoProcessor::default(),
                compression_processor: GpuCompressionProcessor::default(),
                memory_manager: GpuMemoryPoolManager::new(device)
                    .unwrap_or_else(|_| panic!("Failed to create fallback GPU memory manager")),
            }
        })
    }
}

/// Comprehensive GPU capabilities information
#[derive(Debug, Clone)]
pub struct GpuCapabilities {
    /// GPU backend type
    pub backend: GpuBackend,
    /// Total memory in gigabytes
    pub memory_gb: f64,
    /// Number of compute units
    pub compute_units: usize,
    /// Whether FP64 operations are supported
    pub supports_fp64: bool,
    /// Whether FP16 operations are supported
    pub supports_fp16: bool,
    /// Compression throughput in GB/s
    pub compression_throughput_gbps: f64,
    /// Number of memory pools
    pub memory_pools: usize,
    /// Total size of all pools
    pub total_pool_size: usize,
    /// Overall performance score (0.0-1.0)
    pub performance_score: f64,
}

impl GpuCapabilities {
    /// Check if GPU is suitable for high-performance computing
    pub fn is_hpc_capable(&self) -> bool {
        self.memory_gb >= 4.0
            && self.compute_units >= 16
            && self.supports_fp64
            && self.performance_score >= 0.7
    }

    /// Check if GPU is suitable for AI/ML workloads
    pub fn is_ml_capable(&self) -> bool {
        self.memory_gb >= 6.0
            && self.compute_units >= 32
            && (self.supports_fp16 || self.supports_fp64)
            && self.performance_score >= 0.6
    }

    /// Get recommended workload types for this GPU
    pub fn get_recommended_workloads(&self) -> Vec<GpuWorkloadType> {
        let mut workloads = Vec::new();

        if self.is_ml_capable() {
            workloads.push(GpuWorkloadType::MachineLearning);
        }

        if self.memory_gb >= 2.0 && self.compute_units >= 8 {
            workloads.push(GpuWorkloadType::ImageProcessing);
        }

        if self.compression_throughput_gbps >= 1.0 {
            workloads.push(GpuWorkloadType::Compression);
        }

        workloads.push(GpuWorkloadType::GeneralCompute);
        workloads
    }
}

/// Unified GPU performance statistics
#[derive(Debug, Clone)]
pub struct UnifiedGpuStats {
    /// GPU backend type
    pub backend: GpuBackend,
    /// Compression performance statistics
    pub compression_stats: CompressionStats,
    /// Memory pool statistics
    pub memory_stats: GlobalPoolStats,
    /// Overall efficiency score (0.0-1.0)
    pub overall_efficiency: f64,
}

/// Maintenance operation report
#[derive(Debug)]
pub struct MaintenanceReport {
    /// Number of buffers freed during maintenance
    pub freed_buffers: usize,
    /// Timestamp of maintenance operation
    pub timestamp: std::time::Instant,
}

/// Convenience functions for common GPU operations
pub mod utils {
    use super::*;

    /// Check if any GPU backend is available
    pub fn is_gpu_available() -> bool {
        GpuIoProcessor::list_available_backends()
            .iter()
            .any(|&backend| backend != GpuBackend::Cpu)
    }

    /// Get the best available GPU backend for a workload
    pub fn get_best_backend_for_workload(workload: GpuWorkloadType) -> Result<GpuBackend> {
        GpuIoProcessor::get_optimal_backend_for_workload(workload)
    }

    /// Create optimized GPU processor for specific workload
    pub fn create_optimized_processor(workload: GpuWorkloadType) -> Result<UnifiedGpuProcessor> {
        let backend = get_best_backend_for_workload(workload)?;
        let mut processor = UnifiedGpuProcessor::with_backend(backend)?;
        processor.optimize_for_workload(workload)?;
        Ok(processor)
    }

    /// Benchmark GPU performance for different operations
    pub fn benchmark_gpu_performance(processor: &UnifiedGpuProcessor) -> GpuBenchmarkResults {
        let capabilities = processor
            .get_capabilities()
            .unwrap_or_else(|_| GpuCapabilities {
                backend: GpuBackend::Cpu,
                memory_gb: 1.0,
                compute_units: 1,
                supports_fp64: false,
                supports_fp16: false,
                compression_throughput_gbps: 0.1,
                memory_pools: 1,
                total_pool_size: 1024 * 1024,
                performance_score: 0.1,
            });

        GpuBenchmarkResults {
            backend: capabilities.backend,
            memory_bandwidth_gbps: capabilities.memory_gb * 0.8, // Estimate 80% peak
            compute_throughput: capabilities.compute_units as f64 * 100.0, // Arbitrary units
            compression_throughput: capabilities.compression_throughput_gbps,
            overall_score: capabilities.performance_score,
        }
    }
}

/// GPU benchmark results
#[derive(Debug, Clone)]
pub struct GpuBenchmarkResults {
    /// GPU backend type
    pub backend: GpuBackend,
    /// Memory bandwidth in GB/s
    pub memory_bandwidth_gbps: f64,
    /// Compute throughput (arbitrary units)
    pub compute_throughput: f64,
    /// Compression throughput in GB/s
    pub compression_throughput: f64,
    /// Overall performance score (0.0-1.0)
    pub overall_score: f64,
}

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

    #[test]
    fn test_unified_processor_creation() {
        // Should construct successfully and select any supported backend (CPU or GPU)
        let processor = UnifiedGpuProcessor::default();
        let backend = processor.backend();
        match backend {
            GpuBackend::Cpu
            | GpuBackend::Metal
            | GpuBackend::OpenCL
            | GpuBackend::Cuda
            | GpuBackend::Rocm
            | GpuBackend::Wgpu => {}
        }
    }

    #[test]
    fn test_gpu_availability_check() {
        // Should always return true due to CPU fallback
        // Test passes if it doesn't panic (CPU is always available)
    }

    #[test]
    fn test_gpu_capabilities() {
        let processor = UnifiedGpuProcessor::default();
        let capabilities = processor.get_capabilities();
        assert!(capabilities.is_ok());

        let caps = capabilities.expect("Operation failed");
        assert!(caps.memory_gb > 0.0);
        assert!(caps.compute_units > 0);
    }

    #[test]
    fn test_workload_optimization() {
        let mut processor = UnifiedGpuProcessor::default();
        let result = processor.optimize_for_workload(GpuWorkloadType::MachineLearning);
        assert!(result.is_ok());
    }

    #[test]
    fn test_performance_benchmarking() {
        let processor = UnifiedGpuProcessor::default();
        let benchmark = utils::benchmark_gpu_performance(&processor);
        assert!(benchmark.overall_score >= 0.0);
        assert!(benchmark.memory_bandwidth_gbps >= 0.0);
    }

    #[test]
    fn test_maintenance_operations() {
        let mut processor = UnifiedGpuProcessor::default();
        let report = processor.maintenance();
        assert!(report.is_ok());
    }
}