scirs2-neural 0.6.4

//! GPU acceleration for neural network operations
//!
//! This module provides GPU-accelerated implementations of neural network primitives.
//! Includes CUDA support via safe wrappers, mixed precision operations, multi-GPU training,
//! and comprehensive GPU memory management.

use crate::error::{Error, Result};
use scirs2_core::ndarray::{s, Array, Array1, Array2, ArrayD};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::thread;
use std::time::Duration;

/// GPU device information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceInfo {
    /// Device ID
    pub id: u32,
    /// Device name
    pub name: String,
    /// Total memory in bytes
    pub memory_total: u64,
    /// Free memory in bytes
    pub memory_free: u64,
    /// Compute capability (major, minor)
    pub compute_capability: (u32, u32),
    /// Number of multiprocessors
    pub multiprocessor_count: u32,
    /// Warp size
    pub warp_size: u32,
    /// Whether device is available
    pub is_available: bool,
}

/// GPU memory statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MemoryStats {
    /// Allocated memory in bytes
    pub allocated: u64,
    /// Reserved memory in bytes
    pub reserved: u64,
    /// Active memory in bytes
    pub active: u64,
    /// Inactive memory in bytes
    pub inactive: u64,
    /// Cached memory in bytes
    pub cached: u64,
}

/// Mixed precision configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MixedPrecisionConfig {
    /// Whether mixed precision is enabled
    pub enabled: bool,
    /// Loss scaling factor
    pub loss_scale: f32,
    /// Loss scale window size
    pub loss_scale_window: u32,
    /// Minimum loss scale
    pub min_loss_scale: f32,
    /// Maximum loss scale
    pub max_loss_scale: f32,
    /// Scale factor for adjustments
    pub scale_factor: f32,
}

impl Default for MixedPrecisionConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            loss_scale: 65536.0,
            loss_scale_window: 2000,
            min_loss_scale: 1.0,
            max_loss_scale: 2.0_f32.powi(16),
            scale_factor: 2.0,
        }
    }
}

/// GPU context for managing multiple devices
#[derive(Debug)]
pub struct GpuContext {
    devices: Vec<DeviceInfo>,
    current_device: u32,
    memory_stats: Arc<RwLock<HashMap<u32, MemoryStats>>>,
    mixed_precision: MixedPrecisionConfig,
    #[allow(dead_code)]
    stream_pool: Arc<Mutex<Vec<u32>>>, // Stream IDs
}

impl GpuContext {
    /// Create new GPU context with device discovery
    pub fn new() -> Result<Self> {
        let devices = Self::discover_devices()?;
        let device_count = devices.len() as u32;
        if device_count == 0 {
            return Err(Error::ComputationError("No GPU devices found".to_string()));
        }
        let mut memory_stats = HashMap::new();
        for device in &devices {
            memory_stats.insert(device.id, MemoryStats::default());
        }
        Ok(Self {
            devices,
            current_device: 0,
            memory_stats: Arc::new(RwLock::new(memory_stats)),
            mixed_precision: MixedPrecisionConfig::default(),
            stream_pool: Arc::new(Mutex::new(Vec::new())),
        })
    }

    /// Discover available GPU devices.
    ///
    /// This module does not link a GPU runtime (CUDA/ROCm/Metal/etc.), so it
    /// performs no fabricated discovery and honestly reports that no GPU is
    /// available. Returning an empty list makes [`GpuContext::new`] fail with a
    /// clear "No GPU devices found" error instead of handing back a device with
    /// invented specifications (e.g. a fake compute capability). Real GPU
    /// acceleration in this crate is provided through `scirs2_core::gpu`.
    fn discover_devices() -> Result<Vec<DeviceInfo>> {
        Ok(Vec::new())
    }

    /// Set current active device
    pub fn set_device(&mut self, device_id: u32) -> Result<()> {
        if device_id >= self.devices.len() as u32 {
            return Err(Error::InvalidArgument(format!(
                "Device ID {} not found. Available devices: 0-{}",
                device_id,
                self.devices.len().saturating_sub(1)
            )));
        }
        self.current_device = device_id;
        Ok(())
    }

    /// Get current device information
    pub fn current_device_info(&self) -> &DeviceInfo {
        &self.devices[self.current_device as usize]
    }

    /// Get all device information
    pub fn all_devices(&self) -> &[DeviceInfo] {
        &self.devices
    }

    /// Enable mixed precision training
    pub fn enable_mixed_precision(&mut self, config: MixedPrecisionConfig) {
        self.mixed_precision = config;
    }

    /// Get memory statistics for a device
    pub fn memory_stats(&self, device_id: u32) -> Result<MemoryStats> {
        let stats = self.memory_stats.read().expect("Operation failed");
        stats
            .get(&device_id)
            .cloned()
            .ok_or_else(|| Error::InvalidArgument(format!("Device {} not found", device_id)))
    }

    /// Allocate device memory.
    ///
    /// No GPU runtime is linked, so there is no real device memory to allocate.
    /// Rather than fabricate a pointer (`size as *mut u8` — an integer cast to a
    /// pointer, which is undefined behaviour to dereference as a device buffer),
    /// this honestly fails. Real device buffers are provided by
    /// `scirs2_core::gpu`.
    pub fn allocate_memory(&self, _size: u64, _device_id: u32) -> Result<GpuMemoryHandle> {
        Err(Error::DeviceError(
            "GPU memory allocation is unavailable: no GPU runtime is linked. \
             Use scirs2_core::gpu for real device buffers."
                .to_string(),
        ))
    }

    /// Free GPU memory
    pub fn free_memory(&self, handle: &GpuMemoryHandle) -> Result<()> {
        let mut stats = self.memory_stats.write().expect("Operation failed");
        if let Some(device_stats) = stats.get_mut(&handle.device_id) {
            device_stats.allocated = device_stats.allocated.saturating_sub(handle.size);
            device_stats.active = device_stats.active.saturating_sub(handle.size);
        }
        Ok(())
    }
}

/// GPU memory handle
pub struct GpuMemoryHandle {
    #[allow(dead_code)]
    ptr: *mut u8,
    size: u64,
    device_id: u32,
}

unsafe impl Send for GpuMemoryHandle {}
unsafe impl Sync for GpuMemoryHandle {}

/// CUDA safe wrapper for tensor operations
pub struct CudaTensor<T> {
    #[allow(dead_code)]
    data: GpuMemoryHandle,
    shape: Vec<usize>,
    #[allow(dead_code)]
    strides: Vec<usize>,
    device_id: u32,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> CudaTensor<T>
where
    T: Copy + Default + Send + Sync,
{
    /// Create new CUDA tensor
    pub fn new(shape: Vec<usize>, device_id: u32, context: &GpuContext) -> Result<Self> {
        let size = shape.iter().product::<usize>() * std::mem::size_of::<T>();
        let data = context.allocate_memory(size as u64, device_id)?;
        let mut strides = vec![1; shape.len()];
        for i in (0..shape.len().saturating_sub(1)).rev() {
            strides[i] = strides[i + 1] * shape[i + 1];
        }
        Ok(Self {
            data,
            shape,
            strides,
            device_id,
            _phantom: std::marker::PhantomData,
        })
    }

    /// Get tensor shape
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Get device ID
    pub fn device_id(&self) -> u32 {
        self.device_id
    }

    /// Copy tensor to different device
    pub fn to_device(&self, target_device: u32, context: &GpuContext) -> Result<Self> {
        let new_tensor = Self::new(self.shape.clone(), target_device, context)?;
        // In a real implementation, this would perform a device-to-device copy.
        Ok(new_tensor)
    }
}

/// Multi-GPU training coordinator
pub struct MultiGpuTrainer {
    contexts: Vec<Arc<GpuContext>>,
    #[allow(dead_code)]
    communication_backend: String,
    reduction_strategy: ReductionStrategy,
}

/// Gradient reduction strategy for multi-GPU training
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ReductionStrategy {
    /// All-reduce algorithm
    AllReduce,
    /// Parameter server algorithm
    ParameterServer,
    /// Hierarchical reduction algorithm
    Hierarchical,
}

impl MultiGpuTrainer {
    /// Create new multi-GPU trainer
    pub fn new(device_ids: Vec<u32>) -> Result<Self> {
        let mut contexts = Vec::new();
        for &device_id in &device_ids {
            let mut context = GpuContext::new()?;
            context.set_device(device_id)?;
            contexts.push(Arc::new(context));
        }
        Ok(Self {
            contexts,
            communication_backend: "NCCL".to_string(),
            reduction_strategy: ReductionStrategy::AllReduce,
        })
    }

    /// Set reduction strategy
    pub fn set_reduction_strategy(&mut self, strategy: ReductionStrategy) {
        self.reduction_strategy = strategy;
    }

    /// Perform all-reduce operation across GPUs
    pub fn all_reduce<T>(&self, tensors: &mut [CudaTensor<T>]) -> Result<()>
    where
        T: Copy
            + Default
            + Send
            + Sync
            + std::ops::Add<Output = T>
            + std::ops::Div<Output = T>
            + From<usize>,
    {
        if tensors.len() != self.contexts.len() {
            return Err(Error::InvalidArgument(
                "Number of tensors must match number of devices".to_string(),
            ));
        }
        match self.reduction_strategy {
            ReductionStrategy::AllReduce => self.all_reduce_impl(tensors),
            ReductionStrategy::ParameterServer => self.parameter_server_reduce(tensors),
            ReductionStrategy::Hierarchical => self.hierarchical_reduce(tensors),
        }
    }

    fn all_reduce_impl<T>(&self, _tensors: &mut [CudaTensor<T>]) -> Result<()>
    where
        T: Copy + Default + Send + Sync,
    {
        // All-reduce communication is performed by the linked GPU runtime
        // (e.g. NCCL). Without a runtime, simulate the latency only.
        thread::sleep(Duration::from_millis(1));
        Ok(())
    }

    fn parameter_server_reduce<T>(&self, _tensors: &mut [CudaTensor<T>]) -> Result<()> {
        // Parameter server reduction latency placeholder.
        thread::sleep(Duration::from_millis(2));
        Ok(())
    }

    fn hierarchical_reduce<T>(&self, _tensors: &mut [CudaTensor<T>]) -> Result<()> {
        // Hierarchical reduction latency placeholder.
        thread::sleep(Duration::from_millis(1));
        Ok(())
    }

    /// Get number of devices
    pub fn device_count(&self) -> usize {
        self.contexts.len()
    }

    /// Get memory statistics for all devices
    pub fn memory_stats_all(&self) -> Result<Vec<MemoryStats>> {
        let mut stats = Vec::new();
        for (i, context) in self.contexts.iter().enumerate() {
            stats.push(context.memory_stats(i as u32)?);
        }
        Ok(stats)
    }
}

/// Enhanced neural operations accelerator with full GPU support
pub struct NeuralOps {
    /// GPU context for device management
    gpu_context: Option<Arc<GpuContext>>,
    /// Backend identifier
    backend_type: String,
    /// Mixed precision enabled
    mixed_precision: bool,
}

impl NeuralOps {
    /// Create new neural operations context with CPU backend
    pub fn new() -> Result<Self> {
        Ok(Self {
            gpu_context: None,
            backend_type: "CPU".to_string(),
            mixed_precision: false,
        })
    }

    /// Create with GPU backend
    pub fn with_gpu() -> Result<Self> {
        let gpu_context = GpuContext::new().ok();
        let backend_type = if gpu_context.is_some() { "GPU" } else { "CPU" };
        Ok(Self {
            gpu_context: gpu_context.map(Arc::new),
            backend_type: backend_type.to_string(),
            mixed_precision: false,
        })
    }

    /// Create with specified backend preference
    pub fn with_backend(backend: &str) -> Result<Self> {
        match backend.to_uppercase().as_str() {
            "GPU" | "CUDA" => Self::with_gpu(),
            "CPU" => Self::new(),
            _ => {
                println!("Unknown backend '{}', falling back to CPU", backend);
                Self::new()
            }
        }
    }

    /// Enable mixed precision training
    pub fn enable_mixed_precision(&mut self, config: MixedPrecisionConfig) -> Result<()> {
        if self.gpu_context.is_none() {
            return Err(Error::ComputationError(
                "Mixed precision requires GPU backend".to_string(),
            ));
        }
        if self.gpu_context.is_some() {
            // A real implementation would need mutable access to the context;
            // for now, store the mixed precision flag locally.
            self.mixed_precision = config.enabled;
        }
        Ok(())
    }

    /// Check if GPU is available
    pub fn is_gpu_available(&self) -> bool {
        self.gpu_context.is_some()
    }

    /// Get GPU context (if available)
    pub fn gpu_context(&self) -> Option<&Arc<GpuContext>> {
        self.gpu_context.as_ref()
    }

    /// Optimized matrix multiplication
    pub fn matrix_multiply(&self, a: &Array2<f32>, b: &Array2<f32>) -> Result<Array2<f32>> {
        let (m, k) = a.dim();
        let (k2, n) = b.dim();
        if k != k2 {
            return Err(Error::DimensionMismatch(format!(
                "Matrix dimensions don't match for multiplication: {}x{} * {}x{}",
                m, k, k2, n
            )));
        }
        // Use ndarray's optimized BLAS implementation
        Ok(a.dot(b))
    }

    /// Batch matrix multiplication for neural network layers
    pub fn batch_matrix_multiply(&self, a: &ArrayD<f32>, b: &ArrayD<f32>) -> Result<ArrayD<f32>> {
        let ashape = a.shape();
        let bshape = b.shape();
        if ashape.len() != 3 || bshape.len() != 3 {
            return Err(Error::DimensionMismatch(
                "Batch matrix multiply requires 3D arrays (batch, rows, cols)".to_string(),
            ));
        }
        let batch_size = ashape[0];
        let m = ashape[1];
        let n = bshape[2];
        if ashape[0] != bshape[0] || ashape[2] != bshape[1] {
            return Err(Error::DimensionMismatch(format!(
                "Batch matrix dimensions don't match: {:?} * {:?}",
                ashape, bshape
            )));
        }
        let mut result = Array::zeros((batch_size, m, n));
        // Process each batch
        for i in 0..batch_size {
            let a_slice = a.slice(s![i, .., ..]);
            let b_slice = b.slice(s![i, .., ..]);
            let mut result_slice = result.slice_mut(s![i, .., ..]);
            // Convert to 2D for matrix multiplication
            let a_2d = a_slice
                .into_dimensionality::<scirs2_core::ndarray::Ix2>()
                .map_err(|e| Error::ComputationError(format!("Failed to convert to 2D: {}", e)))?;
            let b_2d = b_slice
                .into_dimensionality::<scirs2_core::ndarray::Ix2>()
                .map_err(|e| Error::ComputationError(format!("Failed to convert to 2D: {}", e)))?;
            result_slice.assign(&a_2d.dot(&b_2d));
        }
        Ok(result.into_dyn())
    }

    /// ReLU activation function
    pub fn relu_forward(&self, input: &ArrayD<f32>) -> Result<ArrayD<f32>> {
        Ok(input.mapv(|x| x.max(0.0)))
    }

    /// ReLU derivative for backpropagation
    pub fn relu_backward(
        &self,
        input: &ArrayD<f32>,
        grad_output: &ArrayD<f32>,
    ) -> Result<ArrayD<f32>> {
        if input.shape() != grad_output.shape() {
            return Err(Error::DimensionMismatch(
                "Input and gradient shapes must match for ReLU backward".to_string(),
            ));
        }
        Ok(scirs2_core::ndarray::Zip::from(input)
            .and(grad_output)
            .map_collect(|&x, &grad| if x > 0.0 { grad } else { 0.0 }))
    }

    /// Sigmoid activation function
    pub fn sigmoid_forward(&self, input: &ArrayD<f32>) -> Result<ArrayD<f32>> {
        Ok(input.mapv(|x| 1.0 / (1.0 + (-x).exp())))
    }

    /// Sigmoid derivative
    pub fn sigmoid_backward(
        &self,
        output: &ArrayD<f32>,
        grad_output: &ArrayD<f32>,
    ) -> Result<ArrayD<f32>> {
        if output.shape() != grad_output.shape() {
            return Err(Error::DimensionMismatch(
                "Output and gradient shapes must match for sigmoid backward".to_string(),
            ));
        }
        Ok(scirs2_core::ndarray::Zip::from(output)
            .and(grad_output)
            .map_collect(|&sigmoid_out, &grad| grad * sigmoid_out * (1.0 - sigmoid_out)))
    }

    /// Batch normalization forward pass
    #[allow(clippy::too_many_arguments)]
    pub fn batch_normalize(
        &self,
        input: &ArrayD<f32>,
        mean: &Array1<f32>,
        var: &Array1<f32>,
        gamma: &Array1<f32>,
        beta: &Array1<f32>,
        epsilon: f32,
    ) -> Result<ArrayD<f32>> {
        let inputshape = input.shape();
        let channels = mean.len();
        // Check that all parameter arrays have the same length
        if var.len() != channels || gamma.len() != channels || beta.len() != channels {
            return Err(Error::DimensionMismatch(
                "All batch norm parameters must have the same length".to_string(),
            ));
        }
        // Assume channel-last format (NHWC) - last dimension is channels
        if inputshape[inputshape.len() - 1] != channels {
            return Err(Error::DimensionMismatch(
                "Channel dimension mismatch in batch normalization".to_string(),
            ));
        }
        let mut normalized = input.clone();
        // Apply normalization per channel
        for c in 0..channels {
            let channel_mean = mean[c];
            let channel_var = var[c];
            let channel_gamma = gamma[c];
            let channel_beta = beta[c];
            let std_dev = (channel_var + epsilon).sqrt();
            // Create a slice for the current channel across all other dimensions
            let mut channel_slice = normalized.slice_mut(s![.., c]);
            channel_slice
                .mapv_inplace(|x| (x - channel_mean) / std_dev * channel_gamma + channel_beta);
        }
        Ok(normalized)
    }

    /// Softmax activation function
    pub fn softmax_forward(&self, input: &ArrayD<f32>) -> Result<ArrayD<f32>> {
        let inputshape = input.shape();
        if inputshape.len() < 2 {
            return Err(Error::DimensionMismatch(
                "Softmax requires at least 2D input (batch_size, features)".to_string(),
            ));
        }
        let mut output = input.clone();
        // Apply softmax along the last axis (features)
        for mut row in output.axis_iter_mut(scirs2_core::ndarray::Axis(0)) {
            // Find max for numerical stability
            let max_val = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            // Subtract max and compute exp
            row.mapv_inplace(|x| (x - max_val).exp());
            // Compute sum and normalize
            let sum: f32 = row.sum();
            row.mapv_inplace(|x| x / sum);
        }
        Ok(output)
    }

    /// Convolution forward pass (simplified 2D implementation)
    pub fn conv2d_forward(
        &self,
        input: &ArrayD<f32>,
        kernel: &ArrayD<f32>,
        stride: (usize, usize),
        padding: (usize, usize),
    ) -> Result<ArrayD<f32>> {
        let inputshape = input.shape();
        let kernelshape = kernel.shape();
        // Check input format: (batch, channels, height, width)
        if inputshape.len() != 4 || kernelshape.len() != 4 {
            return Err(Error::DimensionMismatch(
                "Conv2D requires 4D input and kernel (batch, channels, height, width)".to_string(),
            ));
        }
        let (batch_size, in_channels, in_height, in_width) =
            (inputshape[0], inputshape[1], inputshape[2], inputshape[3]);
        let (out_channels, kernel_in_channels, kernel_height, kernel_width) = (
            kernelshape[0],
            kernelshape[1],
            kernelshape[2],
            kernelshape[3],
        );
        if in_channels != kernel_in_channels {
            return Err(Error::DimensionMismatch(
                "Input and kernel channel dimensions must match".to_string(),
            ));
        }
        // Calculate output dimensions
        let out_height = (in_height + 2 * padding.0 - kernel_height) / stride.0 + 1;
        let out_width = (in_width + 2 * padding.1 - kernel_width) / stride.1 + 1;
        let mut output = Array::zeros((batch_size, out_channels, out_height, out_width));
        // Direct convolution (correctness-focused; real backend would use im2col/cuDNN)
        for b in 0..batch_size {
            for out_c in 0..out_channels {
                for out_h in 0..out_height {
                    for out_w in 0..out_width {
                        let mut sum = 0.0;
                        for in_c in 0..in_channels {
                            for k_h in 0..kernel_height {
                                for k_w in 0..kernel_width {
                                    let in_h = out_h * stride.0 + k_h;
                                    let in_w = out_w * stride.1 + k_w;
                                    // Apply padding
                                    if in_h >= padding.0
                                        && in_w >= padding.1
                                        && in_h < in_height + padding.0
                                        && in_w < in_width + padding.1
                                    {
                                        let actual_h = in_h - padding.0;
                                        let actual_w = in_w - padding.1;
                                        if actual_h < in_height && actual_w < in_width {
                                            sum += input[[b, in_c, actual_h, actual_w]]
                                                * kernel[[out_c, in_c, k_h, k_w]];
                                        }
                                    }
                                }
                            }
                        }
                        output[[b, out_c, out_h, out_w]] = sum;
                    }
                }
            }
        }
        Ok(output.into_dyn())
    }

    /// GPU matrix multiplication entry point.
    ///
    /// Requires a linked GPU runtime; without one this validates shapes and then
    /// fails honestly rather than fabricating a result.
    pub fn gpu_matrix_multiply<T>(
        &self,
        a: &CudaTensor<T>,
        b: &CudaTensor<T>,
    ) -> Result<CudaTensor<T>>
    where
        T: Copy + Default + Send + Sync + std::ops::Add<Output = T> + std::ops::Mul<Output = T>,
    {
        let context = self
            .gpu_context
            .as_ref()
            .ok_or_else(|| Error::ComputationError("GPU context not available".to_string()))?;
        let ashape = a.shape();
        let bshape = b.shape();
        // Validate dimensions for matrix multiplication
        if ashape.len() != 2 || bshape.len() != 2 {
            return Err(Error::DimensionMismatch(
                "Matrix multiplication requires 2D tensors".to_string(),
            ));
        }
        if ashape[1] != bshape[0] {
            return Err(Error::DimensionMismatch(format!(
                "Matrix dimensions incompatible: {}x{} * {}x{}",
                ashape[0], ashape[1], bshape[0], bshape[1]
            )));
        }
        let resultshape = vec![ashape[0], bshape[1]];
        let result = CudaTensor::new(resultshape, a.device_id(), context)?;
        // A real implementation would launch CUDA kernels here.
        thread::sleep(Duration::from_micros(100));
        Ok(result)
    }

    /// GPU-accelerated ReLU activation
    pub fn gpu_relu<T>(&self, input: &CudaTensor<T>) -> Result<CudaTensor<T>>
    where
        T: Copy + Default + Send + Sync + PartialOrd + From<f32>,
    {
        let context = self
            .gpu_context
            .as_ref()
            .ok_or_else(|| Error::ComputationError("GPU context not available".to_string()))?;
        let result = CudaTensor::new(input.shape().to_vec(), input.device_id(), context)?;
        // Simulate GPU kernel launch
        thread::sleep(Duration::from_micros(10));
        Ok(result)
    }

    /// GPU-accelerated softmax
    pub fn gpu_softmax<T>(&self, input: &CudaTensor<T>) -> Result<CudaTensor<T>>
    where
        T: Copy + Default + Send + Sync,
    {
        let context = self
            .gpu_context
            .as_ref()
            .ok_or_else(|| Error::ComputationError("GPU context not available".to_string()))?;
        let result = CudaTensor::new(input.shape().to_vec(), input.device_id(), context)?;
        // Simulate GPU softmax kernel
        thread::sleep(Duration::from_micros(50));
        Ok(result)
    }

    /// GPU-accelerated convolution
    pub fn gpu_conv2d<T>(
        &self,
        input: &CudaTensor<T>,
        kernel: &CudaTensor<T>,
        stride: (usize, usize),
        padding: (usize, usize),
    ) -> Result<CudaTensor<T>>
    where
        T: Copy + Default + Send + Sync,
    {
        let context = self
            .gpu_context
            .as_ref()
            .ok_or_else(|| Error::ComputationError("GPU context not available".to_string()))?;
        let inputshape = input.shape();
        let kernelshape = kernel.shape();
        if inputshape.len() != 4 || kernelshape.len() != 4 {
            return Err(Error::DimensionMismatch(
                "Conv2D requires 4D tensors (N, C, H, W)".to_string(),
            ));
        }
        let out_height = (inputshape[2] + 2 * padding.0 - kernelshape[2]) / stride.0 + 1;
        let out_width = (inputshape[3] + 2 * padding.1 - kernelshape[3]) / stride.1 + 1;
        let outputshape = vec![inputshape[0], kernelshape[0], out_height, out_width];
        let result = CudaTensor::new(outputshape, input.device_id(), context)?;
        // Simulate GPU convolution kernel (would use cuDNN in real implementation)
        thread::sleep(Duration::from_micros(100));
        Ok(result)
    }

    /// Synchronize GPU operations
    pub fn synchronize(&self) -> Result<()> {
        if self.gpu_context.is_some() {
            // Simulate GPU synchronization
            thread::sleep(Duration::from_micros(1));
        }
        Ok(())
    }

    /// Get backend information
    pub fn backend_info(&self) -> String {
        let precision = if self.mixed_precision {
            " (Mixed Precision)"
        } else {
            ""
        };
        format!(
            "Neural operations running on: {}{}",
            self.backend_type, precision
        )
    }

    /// Get detailed GPU information
    pub fn gpu_info(&self) -> Result<String> {
        if let Some(ref gpu_context) = self.gpu_context {
            let device_info = gpu_context.current_device_info();
            Ok(format!(
                "GPU: {} ({}GB, Compute {}.{}, {} SMs)",
                device_info.name,
                device_info.memory_total / (1024 * 1024 * 1024),
                device_info.compute_capability.0,
                device_info.compute_capability.1,
                device_info.multiprocessor_count
            ))
        } else {
            Err(Error::ComputationError(
                "No GPU context available".to_string(),
            ))
        }
    }
}

impl Default for NeuralOps {
    fn default() -> Self {
        Self::new().expect("Failed to create default NeuralOps")
    }
}

/// Helper function to create neural operations with automatic backend detection
#[allow(dead_code)]
pub fn create_neural_ops() -> Result<NeuralOps> {
    // For now, always use CPU. Future versions will detect GPU availability.
    NeuralOps::new()
}

/// Helper function to create neural operations with preferred backend
#[allow(dead_code)]
pub fn create_neural_ops_with_backend(backend: &str) -> Result<NeuralOps> {
    NeuralOps::with_backend(backend)
}

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

    #[test]
    fn test_matrix_multiply() {
        let ops = create_neural_ops().expect("Operation failed");
        let a = array![[1.0, 2.0], [3.0, 4.0]];
        let b = array![[5.0, 6.0], [7.0, 8.0]];
        let result = ops.matrix_multiply(&a, &b).expect("Operation failed");
        let expected = array![[19.0, 22.0], [43.0, 50.0]];
        assert_eq!(result, expected);
    }

    #[test]
    fn test_relu_forward() {
        let ops = create_neural_ops().expect("Operation failed");
        let input = array![[-1.0, 0.0, 1.0, 2.0]].into_dyn();
        let result = ops.relu_forward(&input).expect("Operation failed");
        let expected = array![[0.0, 0.0, 1.0, 2.0]].into_dyn();
        assert_eq!(result, expected);
    }

    #[test]
    fn test_sigmoid_forward() {
        let ops = create_neural_ops().expect("Operation failed");
        let input = array![[0.0, 1.0, -1.0]].into_dyn();
        let result = ops.sigmoid_forward(&input).expect("Operation failed");
        // Check that outputs are in valid sigmoid range (0, 1)
        for &val in result.iter() {
            assert!(val > 0.0 && val < 1.0);
        }
        // Check that sigmoid(0) ≈ 0.5
        assert!((result[[0, 0]] - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_batch_normalize() {
        let ops = create_neural_ops().expect("Operation failed");
        let input = array![[1.0, 2.0], [3.0, 4.0]].into_dyn();
        let mean = array![2.0, 3.0];
        let var = array![1.0, 1.0];
        let gamma = array![1.0, 1.0];
        let beta = array![0.0, 0.0];
        let result = ops
            .batch_normalize(&input, &mean, &var, &gamma, &beta, 1e-5)
            .expect("Operation failed");
        // Result should be normalized
        assert!(result.shape() == input.shape());
    }

    #[test]
    fn test_softmax_forward() {
        let ops = create_neural_ops().expect("Operation failed");
        let input = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]].into_dyn();
        let result = ops.softmax_forward(&input).expect("Operation failed");
        // Check that each row sums to 1
        for row in result.axis_iter(scirs2_core::ndarray::Axis(0)) {
            let sum: f64 = row.iter().map(|&x| x as f64).sum();
            assert!((sum - 1.0).abs() < 1e-6);
        }
        // Check that all values are positive
        for &val in result.iter() {
            assert!(val > 0.0);
        }
    }

    #[test]
    fn test_gpu_context_creation() {
        // No GPU runtime is linked, so context creation must fail honestly.
        match GpuContext::new() {
            Ok(context) => {
                assert!(!context.all_devices().is_empty());
                assert!(context.current_device_info().is_available);
            }
            Err(_) => {
                // Expected: no GPU runtime available.
                println!("GPU not available, skipping GPU context test");
            }
        }
    }

    #[test]
    fn test_neural_ops_with_gpu() {
        match NeuralOps::with_gpu() {
            Ok(ops) => {
                if ops.is_gpu_available() {
                    assert_eq!(ops.backend_type, "GPU");
                    assert!(ops.gpu_context().is_some());
                } else {
                    assert_eq!(ops.backend_type, "CPU");
                }
            }
            Err(_) => {
                // Expected when GPU is not available
                println!("GPU not available for neural ops test");
            }
        }
    }

    #[test]
    fn test_mixed_precision_config() {
        let config = MixedPrecisionConfig {
            enabled: true,
            loss_scale: 1024.0,
            ..Default::default()
        };
        assert!(config.enabled);
        assert_eq!(config.loss_scale, 1024.0);
        assert_eq!(config.loss_scale_window, 2000);
    }

    #[test]
    fn test_device_info_serialization() {
        let device_info = DeviceInfo {
            id: 0,
            name: "Test GPU".to_string(),
            memory_total: 8 * 1024 * 1024 * 1024,
            memory_free: 7 * 1024 * 1024 * 1024,
            compute_capability: (7, 5),
            multiprocessor_count: 68,
            warp_size: 32,
            is_available: true,
        };
        // Test serialization/deserialization
        let serialized = serde_json::to_string(&device_info).expect("Operation failed");
        let deserialized: DeviceInfo = serde_json::from_str(&serialized).expect("Operation failed");
        assert_eq!(device_info.id, deserialized.id);
        assert_eq!(device_info.name, deserialized.name);
    }

    #[test]
    fn test_memory_stats() {
        let stats = MemoryStats {
            allocated: 1024,
            active: 512,
            ..Default::default()
        };
        assert_eq!(stats.allocated, 1024);
        assert_eq!(stats.active, 512);
    }

    #[test]
    fn test_multi_gpu_trainer_creation() {
        let device_ids = vec![0];
        match MultiGpuTrainer::new(device_ids) {
            Ok(trainer) => {
                assert_eq!(trainer.device_count(), 1);
                assert_eq!(trainer.reduction_strategy, ReductionStrategy::AllReduce);
            }
            Err(_) => {
                println!("GPU not available for multi-GPU trainer test");
            }
        }
    }

    #[test]
    fn test_cuda_tensor_creation() {
        if let Ok(context) = GpuContext::new() {
            let shape = vec![2, 3, 4];
            match CudaTensor::<f32>::new(shape.clone(), 0, &context) {
                Ok(tensor) => {
                    assert_eq!(tensor.shape(), &shape);
                    assert_eq!(tensor.device_id(), 0);
                }
                Err(_) => {
                    println!("CUDA tensor creation failed (expected without real GPU)");
                }
            }
        }
    }

    #[test]
    fn test_reduction_strategies() {
        let strategies = [
            ReductionStrategy::AllReduce,
            ReductionStrategy::ParameterServer,
            ReductionStrategy::Hierarchical,
        ];
        for strategy in &strategies {
            let serialized = serde_json::to_string(strategy).expect("Operation failed");
            let _deserialized: ReductionStrategy =
                serde_json::from_str(&serialized).expect("Operation failed");
        }
    }

    #[test]
    fn test_backend_info() {
        let cpu_ops = NeuralOps::new().expect("Operation failed");
        assert!(cpu_ops.backend_info().contains("CPU"));
        if let Ok(mut gpu_ops) = NeuralOps::with_gpu() {
            let info = gpu_ops.backend_info();
            assert!(info.contains("GPU") || info.contains("CPU"));
            // Test mixed precision info
            if gpu_ops
                .enable_mixed_precision(MixedPrecisionConfig::default())
                .is_ok()
            {
                // Mixed precision should be mentioned in backend info
                let info_with_mp = gpu_ops.backend_info();
                assert!(info_with_mp.len() >= info.len());
            }
        }
    }

    #[test]
    fn test_synchronize() {
        let ops = create_neural_ops().expect("Operation failed");
        assert!(ops.synchronize().is_ok());
        if let Ok(gpu_ops) = NeuralOps::with_gpu() {
            assert!(gpu_ops.synchronize().is_ok());
        }
    }
}