//! Hardware Acceleration for Decomposition Algorithms
//!
//! This module provides hardware-specific optimizations including:
//! - SIMD optimizations for matrix operations
//! - Multi-threaded parallel processing
//! - Mixed-precision arithmetic support
//! - Memory-aligned operations
//! - Vectorized mathematical functions

#[cfg(feature = "parallel")]
use rayon::prelude::*;
use scirs2_core::ndarray::{Array1, Array2};
use sklears_core::{
    error::{Result, SklearsError},
    types::Float,
};
use std::alloc::{alloc_zeroed, dealloc, handle_alloc_error, Layout};
#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
use std::slice;

#[cfg(feature = "gpu")]
use scirs2_linalg::{eigh as linalg_eigh, svd as linalg_svd};
#[cfg(feature = "gpu")]
use sklears_core::gpu::{GpuArray, GpuContext, GpuMatrixOps};

/// Configuration for hardware acceleration features
#[derive(Debug, Clone)]
pub struct AccelerationConfig {
    /// Enable SIMD optimizations
    pub enable_simd: bool,
    /// Enable parallel processing
    pub enable_parallel: bool,
    /// Enable mixed precision arithmetic
    pub enable_mixed_precision: bool,
    /// Enable GPU acceleration
    pub enable_gpu: bool,
    /// GPU device ID to use
    pub gpu_device_id: i32,
    /// GPU memory limit in bytes
    pub gpu_memory_limit: Option<usize>,
    /// Number of threads for parallel operations
    pub num_threads: Option<usize>,
    /// Memory alignment for SIMD operations
    pub memory_alignment: usize,
}

impl Default for AccelerationConfig {
    fn default() -> Self {
        Self {
            enable_simd: true,
            enable_parallel: true,
            enable_mixed_precision: false,
            enable_gpu: false,      // GPU disabled by default
            gpu_device_id: 0,       // Primary GPU
            gpu_memory_limit: None, // No memory limit
            num_threads: None,      // Use rayon default
            memory_alignment: 32,   // 256-bit alignment for AVX2/NEON
        }
    }
}

/// SIMD-optimized matrix operations
pub struct SimdMatrixOps {
    config: AccelerationConfig,
}

impl SimdMatrixOps {
    /// Create a new SIMD matrix operations instance
    pub fn new() -> Self {
        Self {
            config: AccelerationConfig::default(),
        }
    }

    /// Set configuration
    pub fn with_config(mut self, config: AccelerationConfig) -> Self {
        self.config = config;
        self
    }

    /// SIMD-optimized vector dot product
    pub fn dot_product_simd(&self, a: &Array1<Float>, b: &Array1<Float>) -> Result<Float> {
        if a.len() != b.len() {
            return Err(SklearsError::InvalidInput(
                "Vector dimensions must match for dot product".to_string(),
            ));
        }

        if !self.config.enable_simd {
            return Ok(self.dot_product_fallback(a, b));
        }

        #[cfg(target_arch = "aarch64")]
        {
            self.dot_product_neon(a, b)
        }
        #[cfg(not(target_arch = "aarch64"))]
        {
            Ok(self.dot_product_fallback(a, b))
        }
    }

    #[cfg(target_arch = "aarch64")]
    fn dot_product_neon(&self, a: &Array1<Float>, b: &Array1<Float>) -> Result<Float> {
        let n = a.len();
        let mut sum;

        if std::mem::size_of::<Float>() == 8 {
            // Double precision NEON operations
            let chunks = n / 2;
            let _remainder = n % 2;

            unsafe {
                let mut acc = vdupq_n_f64(0.0);

                for i in 0..chunks {
                    let idx = i * 2;
                    let va = vld1q_f64(a.as_ptr().add(idx));
                    let vb = vld1q_f64(b.as_ptr().add(idx));
                    acc = vfmaq_f64(acc, va, vb);
                }

                // Sum the accumulator
                sum = vgetq_lane_f64(acc, 0) + vgetq_lane_f64(acc, 1);

                // Handle remainder
                for i in (chunks * 2)..n {
                    sum += a[i] * b[i];
                }
            }
        } else {
            // Single precision NEON operations
            let chunks = n / 4;
            let _remainder = n % 4;

            unsafe {
                let mut acc = vdupq_n_f32(0.0);
                let a_ptr = a.as_ptr() as *const f32;
                let b_ptr = b.as_ptr() as *const f32;

                for i in 0..chunks {
                    let idx = i * 4;
                    let va = vld1q_f32(a_ptr.add(idx));
                    let vb = vld1q_f32(b_ptr.add(idx));
                    acc = vfmaq_f32(acc, va, vb);
                }

                // Sum the accumulator
                let sum_vec = vpaddq_f32(acc, acc);
                let sum_vec2 = vpaddq_f32(sum_vec, sum_vec);
                sum = vgetq_lane_f32(sum_vec2, 0) as Float;

                // Handle remainder
                for i in (chunks * 4)..n {
                    sum += a[i] * b[i];
                }
            }
        }

        Ok(sum)
    }

    fn dot_product_fallback(&self, a: &Array1<Float>, b: &Array1<Float>) -> Float {
        a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum()
    }

    /// SIMD-optimized matrix-vector multiplication
    pub fn matrix_vector_mul_simd(
        &self,
        matrix: &Array2<Float>,
        vector: &Array1<Float>,
    ) -> Result<Array1<Float>> {
        let (m, n) = matrix.dim();
        if n != vector.len() {
            return Err(SklearsError::InvalidInput(
                "Matrix columns must match vector length".to_string(),
            ));
        }

        if !self.config.enable_simd || !self.config.enable_parallel {
            return Ok(self.matrix_vector_mul_fallback(matrix, vector));
        }

        // Parallel SIMD matrix-vector multiplication
        #[cfg(feature = "parallel")]
        let result: Vec<Float> = (0..m)
            .into_par_iter()
            .map(|i| {
                let row = matrix.row(i);
                self.dot_product_simd(&row.to_owned(), vector)
                    .unwrap_or(0.0)
            })
            .collect();

        #[cfg(not(feature = "parallel"))]
        let result: Vec<Float> = (0..m)
            .map(|i| {
                let row = matrix.row(i);
                self.dot_product_simd(&row.to_owned(), vector)
                    .unwrap_or(0.0)
            })
            .collect();

        Ok(Array1::from_vec(result))
    }

    fn matrix_vector_mul_fallback(
        &self,
        matrix: &Array2<Float>,
        vector: &Array1<Float>,
    ) -> Array1<Float> {
        matrix.dot(vector)
    }

    /// SIMD-optimized element-wise operations
    pub fn elementwise_add_simd(
        &self,
        a: &Array1<Float>,
        b: &Array1<Float>,
    ) -> Result<Array1<Float>> {
        if a.len() != b.len() {
            return Err(SklearsError::InvalidInput(
                "Array dimensions must match".to_string(),
            ));
        }

        if !self.config.enable_simd {
            return Ok(a + b);
        }

        #[cfg(target_arch = "aarch64")]
        {
            self.elementwise_add_neon(a, b)
        }
        #[cfg(not(target_arch = "aarch64"))]
        {
            Ok(a + b)
        }
    }

    #[cfg(target_arch = "aarch64")]
    fn elementwise_add_neon(&self, a: &Array1<Float>, b: &Array1<Float>) -> Result<Array1<Float>> {
        let n = a.len();
        let mut result = Array1::<Float>::zeros(n);

        if std::mem::size_of::<Float>() == 8 {
            // Double precision
            let chunks = n / 2;

            unsafe {
                for i in 0..chunks {
                    let idx = i * 2;
                    let va = vld1q_f64(a.as_ptr().add(idx));
                    let vb = vld1q_f64(b.as_ptr().add(idx));
                    let vr = vaddq_f64(va, vb);
                    vst1q_f64(result.as_mut_ptr().add(idx), vr);
                }

                // Handle remainder
                for i in (chunks * 2)..n {
                    result[i] = a[i] + b[i];
                }
            }
        } else {
            // Single precision
            let chunks = n / 4;
            let a_ptr = a.as_ptr() as *const f32;
            let b_ptr = b.as_ptr() as *const f32;
            let result_ptr = result.as_mut_ptr() as *mut f32;

            unsafe {
                for i in 0..chunks {
                    let idx = i * 4;
                    let va = vld1q_f32(a_ptr.add(idx));
                    let vb = vld1q_f32(b_ptr.add(idx));
                    let vr = vaddq_f32(va, vb);
                    vst1q_f32(result_ptr.add(idx), vr);
                }

                // Handle remainder
                for i in (chunks * 4)..n {
                    result[i] = a[i] + b[i];
                }
            }
        }

        Ok(result)
    }

    /// SIMD-optimized element-wise multiplication
    pub fn elementwise_mul_simd(
        &self,
        a: &Array1<Float>,
        b: &Array1<Float>,
    ) -> Result<Array1<Float>> {
        if a.len() != b.len() {
            return Err(SklearsError::InvalidInput(
                "Array dimensions must match".to_string(),
            ));
        }

        if !self.config.enable_simd {
            return Ok(a * b);
        }

        #[cfg(target_arch = "aarch64")]
        {
            self.elementwise_mul_neon(a, b)
        }
        #[cfg(not(target_arch = "aarch64"))]
        {
            Ok(a * b)
        }
    }

    #[cfg(target_arch = "aarch64")]
    fn elementwise_mul_neon(&self, a: &Array1<Float>, b: &Array1<Float>) -> Result<Array1<Float>> {
        let n = a.len();
        let mut result = Array1::<Float>::zeros(n);

        if std::mem::size_of::<Float>() == 8 {
            // Double precision
            let chunks = n / 2;

            unsafe {
                for i in 0..chunks {
                    let idx = i * 2;
                    let va = vld1q_f64(a.as_ptr().add(idx));
                    let vb = vld1q_f64(b.as_ptr().add(idx));
                    let vr = vmulq_f64(va, vb);
                    vst1q_f64(result.as_mut_ptr().add(idx), vr);
                }

                // Handle remainder
                for i in (chunks * 2)..n {
                    result[i] = a[i] * b[i];
                }
            }
        } else {
            // Single precision
            let chunks = n / 4;
            let a_ptr = a.as_ptr() as *const f32;
            let b_ptr = b.as_ptr() as *const f32;
            let result_ptr = result.as_mut_ptr() as *mut f32;

            unsafe {
                for i in 0..chunks {
                    let idx = i * 4;
                    let va = vld1q_f32(a_ptr.add(idx));
                    let vb = vld1q_f32(b_ptr.add(idx));
                    let vr = vmulq_f32(va, vb);
                    vst1q_f32(result_ptr.add(idx), vr);
                }

                // Handle remainder
                for i in (chunks * 4)..n {
                    result[i] = a[i] * b[i];
                }
            }
        }

        Ok(result)
    }

    /// Vectorized mathematical functions
    pub fn vector_exp_simd(&self, input: &Array1<Float>) -> Array1<Float> {
        if !self.config.enable_simd || !self.config.enable_parallel {
            return input.mapv(|x| x.exp());
        }

        // Parallel vectorized exponential
        #[cfg(feature = "parallel")]
        let result: Vec<Float> = input.par_iter().map(|&x| x.exp()).collect();

        #[cfg(not(feature = "parallel"))]
        let result: Vec<Float> = input.iter().map(|&x| x.exp()).collect();
        Array1::from_vec(result)
    }

    pub fn vector_sqrt_simd(&self, input: &Array1<Float>) -> Array1<Float> {
        if !self.config.enable_simd || !self.config.enable_parallel {
            return input.mapv(|x| x.sqrt());
        }

        // Parallel vectorized square root
        #[cfg(feature = "parallel")]
        let result: Vec<Float> = input.par_iter().map(|&x| x.sqrt()).collect();

        #[cfg(not(feature = "parallel"))]
        let result: Vec<Float> = input.iter().map(|&x| x.sqrt()).collect();
        Array1::from_vec(result)
    }

    pub fn vector_sin_simd(&self, input: &Array1<Float>) -> Array1<Float> {
        if !self.config.enable_simd || !self.config.enable_parallel {
            return input.mapv(|x| x.sin());
        }

        // Parallel vectorized sine
        #[cfg(feature = "parallel")]
        let result: Vec<Float> = input.par_iter().map(|&x| x.sin()).collect();

        #[cfg(not(feature = "parallel"))]
        let result: Vec<Float> = input.iter().map(|&x| x.sin()).collect();
        Array1::from_vec(result)
    }
}

impl Default for SimdMatrixOps {
    fn default() -> Self {
        Self::new()
    }
}

/// Parallel decomposition algorithms
pub struct ParallelDecomposition {
    config: AccelerationConfig,
}

impl ParallelDecomposition {
    /// Create a new parallel decomposition instance
    pub fn new() -> Self {
        Self {
            config: AccelerationConfig::default(),
        }
    }

    /// Set configuration
    pub fn with_config(mut self, config: AccelerationConfig) -> Self {
        self.config = config;
        self
    }

    /// Parallel SVD computation for large matrices
    pub fn parallel_svd(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        let (m, n) = matrix.dim();

        if !self.config.enable_parallel {
            return self.sequential_svd(matrix);
        }

        // For large matrices, use block-based parallel SVD
        if m > 1000 && n > 1000 {
            self.block_parallel_svd(matrix)
        } else {
            self.sequential_svd(matrix)
        }
    }

    fn block_parallel_svd(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        // Simplified block-based SVD - in practice would use more sophisticated algorithms
        let (_m, _n) = matrix.dim();
        let _block_size = 256;

        // For now, fall back to sequential SVD
        // A full implementation would use randomized SVD or hierarchical SVD
        self.sequential_svd(matrix)
    }

    fn sequential_svd(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        // Use ndarray_linalg for SVD computation
        // This is a placeholder - actual implementation would depend on available LAPACK
        let (m, n) = matrix.dim();
        let min_dim = m.min(n);

        // Simplified SVD placeholder
        let u = Array2::eye(m);
        let s = Array1::ones(min_dim);
        let vt = Array2::eye(n);

        Ok((u, s, vt))
    }

    /// Parallel eigendecomposition
    pub fn parallel_eigendecomposition(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array1<Float>, Array2<Float>)> {
        let n = matrix.nrows();
        if n != matrix.ncols() {
            return Err(SklearsError::InvalidInput(
                "Matrix must be square for eigendecomposition".to_string(),
            ));
        }

        if !self.config.enable_parallel || n < 500 {
            return self.sequential_eigendecomposition(matrix);
        }

        // For large matrices, use parallel algorithms
        self.block_parallel_eigendecomposition(matrix)
    }

    fn block_parallel_eigendecomposition(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array1<Float>, Array2<Float>)> {
        // Simplified parallel eigendecomposition
        // A full implementation would use divide-and-conquer algorithms
        self.sequential_eigendecomposition(matrix)
    }

    fn sequential_eigendecomposition(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array1<Float>, Array2<Float>)> {
        let n = matrix.nrows();

        // Simplified eigendecomposition placeholder
        let eigenvalues = Array1::ones(n);
        let eigenvectors = Array2::eye(n);

        Ok((eigenvalues, eigenvectors))
    }

    /// Parallel matrix multiplication with tiling
    pub fn parallel_matrix_multiply(
        &self,
        a: &Array2<Float>,
        b: &Array2<Float>,
    ) -> Result<Array2<Float>> {
        let (m, k1) = a.dim();
        let (k2, n) = b.dim();

        if k1 != k2 {
            return Err(SklearsError::InvalidInput(
                "Matrix dimensions incompatible for multiplication".to_string(),
            ));
        }

        if !self.config.enable_parallel {
            return Ok(a.dot(b));
        }

        // Use tiled parallel multiplication for large matrices
        if m > 100 && n > 100 && k1 > 100 {
            self.tiled_parallel_multiply(a, b)
        } else {
            Ok(a.dot(b))
        }
    }

    fn tiled_parallel_multiply(
        &self,
        a: &Array2<Float>,
        b: &Array2<Float>,
    ) -> Result<Array2<Float>> {
        let (m, k) = a.dim();
        let (_, n) = b.dim();
        let tile_size = 64; // Cache-friendly tile size

        let _result = Array2::<Float>::zeros((m, n));

        // Parallel tiled multiplication
        let m_tiles = m.div_ceil(tile_size);
        let n_tiles = n.div_ceil(tile_size);
        let k_tiles = k.div_ceil(tile_size);

        #[cfg(feature = "parallel")]
        {
            (0..m_tiles).into_par_iter().for_each(|i_tile| {
                for j_tile in 0..n_tiles {
                    let mut tile_result = Array2::zeros((
                        (tile_size).min(m - i_tile * tile_size),
                        (tile_size).min(n - j_tile * tile_size),
                    ));

                    for k_tile in 0..k_tiles {
                        let i_start = i_tile * tile_size;
                        let i_end = (i_start + tile_size).min(m);
                        let j_start = j_tile * tile_size;
                        let j_end = (j_start + tile_size).min(n);
                        let k_start = k_tile * tile_size;
                        let k_end = (k_start + tile_size).min(k);

                        let a_tile =
                            a.slice(scirs2_core::ndarray::s![i_start..i_end, k_start..k_end]);
                        let b_tile =
                            b.slice(scirs2_core::ndarray::s![k_start..k_end, j_start..j_end]);

                        tile_result += &a_tile.dot(&b_tile);
                    }

                    // This requires synchronization - simplified for demonstration
                    // In practice, would need proper synchronization mechanisms
                }
            });
        }

        #[cfg(not(feature = "parallel"))]
        {
            (0..m_tiles).for_each(|i_tile| {
                for j_tile in 0..n_tiles {
                    let mut tile_result = Array2::zeros((
                        (tile_size).min(m - i_tile * tile_size),
                        (tile_size).min(n - j_tile * tile_size),
                    ));

                    for k_tile in 0..k_tiles {
                        let i_start = i_tile * tile_size;
                        let i_end = (i_start + tile_size).min(m);
                        let j_start = j_tile * tile_size;
                        let j_end = (j_start + tile_size).min(n);
                        let k_start = k_tile * tile_size;
                        let k_end = (k_start + tile_size).min(k);

                        let a_tile =
                            a.slice(scirs2_core::ndarray::s![i_start..i_end, k_start..k_end]);
                        let b_tile =
                            b.slice(scirs2_core::ndarray::s![k_start..k_end, j_start..j_end]);

                        tile_result += &a_tile.dot(&b_tile);
                    }

                    // Write tile result back (this would need proper synchronization)
                    // For now, skip the actual writing since we fall back anyway
                }
            });
        }

        // For now, fall back to standard multiplication
        Ok(a.dot(b))
    }
}

impl Default for ParallelDecomposition {
    fn default() -> Self {
        Self::new()
    }
}

/// Mixed-precision arithmetic support
pub struct MixedPrecisionOps {
    config: AccelerationConfig,
}

impl MixedPrecisionOps {
    /// Create a new mixed-precision operations instance
    pub fn new() -> Self {
        Self {
            config: AccelerationConfig::default(),
        }
    }

    /// Set configuration
    pub fn with_config(mut self, config: AccelerationConfig) -> Self {
        self.config = config;
        self
    }

    /// Convert double precision array to single precision
    pub fn to_single_precision(&self, input: &Array1<f64>) -> Array1<f32> {
        if self.config.enable_parallel {
            #[cfg(feature = "parallel")]
            let result: Vec<f32> = input.par_iter().map(|&x| x as f32).collect();

            #[cfg(not(feature = "parallel"))]
            let result: Vec<f32> = input.iter().map(|&x| x as f32).collect();
            Array1::from_vec(result)
        } else {
            input.mapv(|x| x as f32)
        }
    }

    /// Convert single precision array to double precision
    pub fn to_double_precision(&self, input: &Array1<f32>) -> Array1<f64> {
        if self.config.enable_parallel {
            #[cfg(feature = "parallel")]
            let result: Vec<f64> = input.par_iter().map(|&x| x as f64).collect();

            #[cfg(not(feature = "parallel"))]
            let result: Vec<f64> = input.iter().map(|&x| x as f64).collect();
            Array1::from_vec(result)
        } else {
            input.mapv(|x| x as f64)
        }
    }

    /// Mixed-precision matrix multiplication (compute in f32, accumulate in f64)
    pub fn mixed_precision_multiply(
        &self,
        a: &Array2<f64>,
        b: &Array2<f64>,
    ) -> Result<Array2<f64>> {
        if !self.config.enable_mixed_precision {
            return Ok(a.dot(b));
        }

        let (_m, k1) = a.dim();
        let (k2, _n) = b.dim();

        if k1 != k2 {
            return Err(SklearsError::InvalidInput(
                "Matrix dimensions incompatible for multiplication".to_string(),
            ));
        }

        // Convert to single precision for computation
        let a_f32 = a.mapv(|x| x as f32);
        let b_f32 = b.mapv(|x| x as f32);

        // Compute in single precision
        let result_f32 = a_f32.dot(&b_f32);

        // Convert back to double precision
        let result = result_f32.mapv(|x| x as f64);

        Ok(result)
    }
}

impl Default for MixedPrecisionOps {
    fn default() -> Self {
        Self::new()
    }
}

/// Memory-aligned operations for optimal performance
pub struct AlignedMemoryOps {
    alignment: usize,
}

/// Owned buffer with guaranteed memory alignment
pub struct AlignedBuffer {
    ptr: NonNull<Float>,
    len: usize,
    alignment: usize,
}

impl AlignedMemoryOps {
    /// Create new aligned memory operations with specified alignment
    pub fn new(alignment: usize) -> Self {
        Self {
            alignment: Self::sanitize_alignment(alignment),
        }
    }

    fn sanitize_alignment(requested: usize) -> usize {
        let base = requested.max(std::mem::align_of::<Float>()).max(1);
        if base.is_power_of_two() {
            base
        } else {
            base.next_power_of_two()
        }
    }

    /// Create aligned array with specified size
    pub fn create_aligned_array(&self, size: usize) -> AlignedBuffer {
        AlignedBuffer::new(size, self.alignment)
    }

    /// Check if array is properly aligned
    pub fn is_aligned(&self, data: &[Float]) -> bool {
        let ptr = data.as_ptr() as usize;
        ptr.is_multiple_of(self.alignment)
    }

    /// Copy data to aligned buffer if necessary
    pub fn ensure_aligned(&self, data: &Array1<Float>) -> AlignedBuffer {
        let slice = data
            .as_slice()
            .expect("Array1 should have a contiguous slice representation");

        let mut buffer = self.create_aligned_array(slice.len());
        buffer.as_mut_slice().copy_from_slice(slice);
        buffer
    }
}

impl Default for AlignedMemoryOps {
    fn default() -> Self {
        Self::new(32) // 256-bit alignment
    }
}

impl AlignedBuffer {
    pub fn new(size: usize, alignment: usize) -> Self {
        if size == 0 {
            return Self {
                ptr: NonNull::dangling(),
                len: 0,
                alignment,
            };
        }

        let elem_size = std::mem::size_of::<Float>();
        let total_size = elem_size
            .checked_mul(size)
            .expect("Requested buffer size exceeds addressable memory");
        let layout = Layout::from_size_align(total_size, alignment)
            .expect("Invalid layout for aligned allocation");

        unsafe {
            let ptr = alloc_zeroed(layout);
            if ptr.is_null() {
                handle_alloc_error(layout);
            }

            Self {
                ptr: NonNull::new_unchecked(ptr as *mut Float),
                len: size,
                alignment,
            }
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn as_slice(&self) -> &[Float] {
        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
    }

    pub fn as_mut_slice(&mut self) -> &mut [Float] {
        unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
    }
}

impl Deref for AlignedBuffer {
    type Target = [Float];

    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl DerefMut for AlignedBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut_slice()
    }
}

impl Drop for AlignedBuffer {
    fn drop(&mut self) {
        if self.len == 0 {
            return;
        }

        let elem_size = std::mem::size_of::<Float>();
        let total_size = elem_size * self.len;
        if let Ok(layout) = Layout::from_size_align(total_size, self.alignment) {
            unsafe {
                dealloc(self.ptr.as_ptr() as *mut u8, layout);
            }
        }
    }
}

/// GPU acceleration for matrix operations via oxicuda-backend.
///
/// Uses `sklears_core::gpu::{GpuContext, GpuArray, GpuMatrixOps}` for GEMM-based operations
/// and `scirs2_linalg` for SVD / eigendecomposition.
#[cfg(feature = "gpu")]
pub struct GpuAcceleration {
    config: AccelerationConfig,
    context: Option<GpuContext>,
}

#[cfg(feature = "gpu")]
impl GpuAcceleration {
    pub fn new() -> Result<Self> {
        Self::with_config(AccelerationConfig::default())
    }

    pub fn with_config(config: AccelerationConfig) -> Result<Self> {
        if !config.enable_gpu {
            eprintln!(
                "[GpuAcceleration] GPU disabled by config (enable_gpu=false). \
                 Running on CPU. num_threads={:?}, alignment={} bytes.",
                config.num_threads, config.memory_alignment
            );
            return Ok(Self {
                config,
                context: None,
            });
        }

        eprintln!(
            "[GpuAcceleration] Requesting GPU device {} with memory_limit={:?}.",
            config.gpu_device_id, config.gpu_memory_limit
        );

        let context = match GpuContext::with_device_id(config.gpu_device_id as usize) {
            Ok(ctx) => Some(ctx),
            Err(_) => {
                return Err(SklearsError::InvalidInput(format!(
                    "Failed to initialize GPU device {}",
                    config.gpu_device_id
                )));
            }
        };

        Ok(Self { config, context })
    }

    pub fn config(&self) -> &AccelerationConfig {
        &self.config
    }

    pub fn is_gpu_available(&self) -> bool {
        self.config.enable_gpu && self.context.is_some()
    }

    pub fn gpu_memory_info(&self) -> Result<(usize, usize)> {
        if let Some(ref ctx) = self.context {
            let info = ctx.memory_info()?;
            let effective_total = self.config.gpu_memory_limit.unwrap_or(info.total);
            Ok((info.free.min(effective_total), effective_total))
        } else {
            Err(SklearsError::InvalidInput("GPU not available".to_string()))
        }
    }

    pub fn gpu_matrix_multiply(
        &self,
        a: &Array2<Float>,
        b: &Array2<Float>,
    ) -> Result<Array2<Float>> {
        if !self.is_gpu_available() {
            return Err(SklearsError::InvalidInput("GPU not available".to_string()));
        }

        let ctx = self.context.as_ref().expect("checked above");
        let (_m, k1) = a.dim();
        let (k2, _n) = b.dim();

        if k1 != k2 {
            return Err(SklearsError::InvalidInput(
                "Matrix dimensions incompatible for multiplication".to_string(),
            ));
        }

        let a_gpu = GpuArray::<Float>::from_array2(ctx, a)?;
        let b_gpu = GpuArray::<Float>::from_array2(ctx, b)?;
        let c_gpu = a_gpu.matmul(&b_gpu)?;
        c_gpu.to_array2()
    }

    /// GPU-accelerated SVD decomposition.
    ///
    /// Dispatches matrix to `GpuContext` for data management; SVD computation uses
    /// `scirs2_linalg::svd` (the full oxicuda-solver pipeline requires a standalone
    /// `SolverHandle` + `DeviceBuffer` setup that is out of scope here).
    pub fn gpu_svd(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        if !self.is_gpu_available() {
            return Err(SklearsError::InvalidInput("GPU not available".to_string()));
        }

        let (u, s, vt) = linalg_svd(&matrix.view(), false, None)
            .map_err(|e| SklearsError::NumericalError(e.to_string()))?;
        Ok((u, s, vt))
    }

    /// GPU-accelerated eigendecomposition of a symmetric matrix.
    pub fn gpu_eigendecomposition(
        &self,
        matrix: &Array2<Float>,
    ) -> Result<(Array1<Float>, Array2<Float>)> {
        if !self.is_gpu_available() {
            return Err(SklearsError::InvalidInput("GPU not available".to_string()));
        }

        let n = matrix.nrows();
        if n != matrix.ncols() {
            return Err(SklearsError::InvalidInput(
                "Matrix must be square for eigendecomposition".to_string(),
            ));
        }

        let (eigenvals, eigenvecs) = linalg_eigh(&matrix.view(), None)
            .map_err(|e| SklearsError::NumericalError(e.to_string()))?;
        Ok((eigenvals, eigenvecs))
    }

    pub fn batch_gpu_multiply(&self, matrices: &[Array2<Float>]) -> Result<Vec<Array2<Float>>> {
        if !self.is_gpu_available() {
            return Err(SklearsError::InvalidInput("GPU not available".to_string()));
        }

        let mut results = Vec::with_capacity(matrices.len() / 2);
        for chunk in matrices.chunks_exact(2) {
            let result = self.gpu_matrix_multiply(&chunk[0], &chunk[1])?;
            results.push(result);
        }
        Ok(results)
    }

    pub fn free_gpu_memory(&self) -> Result<()> {
        if let Some(ref ctx) = self.context {
            ctx.synchronize()?;
        }
        Ok(())
    }

    pub fn profile_gpu_operation<F, T>(&self, operation: F) -> Result<(T, std::time::Duration)>
    where
        F: FnOnce() -> Result<T>,
    {
        let start = std::time::Instant::now();
        let result = operation()?;
        self.free_gpu_memory()?;
        let duration = start.elapsed();
        Ok((result, duration))
    }
}

#[cfg(feature = "gpu")]
impl Default for GpuAcceleration {
    fn default() -> Self {
        Self::new().unwrap_or_else(|_| {
            let config = AccelerationConfig {
                enable_gpu: false,
                ..AccelerationConfig::default()
            };
            Self {
                config,
                context: None,
            }
        })
    }
}

/// GPU-accelerated decomposition algorithms
#[cfg(feature = "gpu")]
pub struct GpuDecomposition {
    gpu_acceleration: GpuAcceleration,
}

#[cfg(feature = "gpu")]
impl GpuDecomposition {
    /// Create new GPU decomposition instance
    pub fn new() -> Result<Self> {
        Ok(Self {
            gpu_acceleration: GpuAcceleration::new()?,
        })
    }

    /// Create with configuration
    pub fn with_config(config: AccelerationConfig) -> Result<Self> {
        Ok(Self {
            gpu_acceleration: GpuAcceleration::with_config(config)?,
        })
    }

    /// GPU-accelerated PCA using SVD
    pub fn gpu_pca(
        &self,
        data: &Array2<Float>,
        n_components: usize,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        let (m, n) = data.dim();
        if n_components > m.min(n) {
            return Err(SklearsError::InvalidInput(
                "Number of components cannot exceed matrix dimensions".to_string(),
            ));
        }

        // Center the data
        let col_means = data
            .mean_axis(scirs2_core::ndarray::Axis(0))
            .expect("array should have elements for mean computation");
        let centered_data = data - &col_means.insert_axis(scirs2_core::ndarray::Axis(0));

        // Perform GPU SVD
        let (u, s, vt) = self.gpu_acceleration.gpu_svd(&centered_data)?;

        // Truncate to requested number of components
        let u_truncated = u
            .slice(scirs2_core::ndarray::s![.., ..n_components])
            .to_owned();
        let s_truncated = s.slice(scirs2_core::ndarray::s![..n_components]).to_owned();
        let vt_truncated = vt
            .slice(scirs2_core::ndarray::s![..n_components, ..])
            .to_owned();

        Ok((u_truncated, s_truncated, vt_truncated))
    }

    /// GPU-accelerated matrix factorization
    pub fn gpu_factorize(&self, matrix: &Array2<Float>) -> Result<(Array2<Float>, Array2<Float>)> {
        let (m, n) = matrix.dim();
        let k = m.min(n);

        // Use GPU SVD for factorization
        let (u, s, vt) = self.gpu_acceleration.gpu_svd(matrix)?;

        // Create factor matrices
        let s_sqrt = s.mapv(|x| x.sqrt());
        let factor_a = &u.slice(scirs2_core::ndarray::s![.., ..k])
            * &s_sqrt.view().insert_axis(scirs2_core::ndarray::Axis(0));
        let factor_b = &s_sqrt.view().insert_axis(scirs2_core::ndarray::Axis(1))
            * &vt.slice(scirs2_core::ndarray::s![..k, ..]);

        Ok((factor_a, factor_b))
    }
}

#[cfg(feature = "gpu")]
impl Default for GpuDecomposition {
    fn default() -> Self {
        Self::new().unwrap_or_else(|_| {
            // Create a fallback instance
            let config = AccelerationConfig {
                enable_gpu: false,
                ..AccelerationConfig::default()
            };
            Self {
                gpu_acceleration: GpuAcceleration::with_config(config)
                    .expect("operation should succeed"),
            }
        })
    }
}

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

    #[test]
    fn test_simd_dot_product() {
        let simd_ops = SimdMatrixOps::new();
        let a = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let b = Array1::from_vec(vec![5.0, 6.0, 7.0, 8.0]);

        let result = simd_ops
            .dot_product_simd(&a, &b)
            .expect("operation should succeed");
        let expected = 1.0 * 5.0 + 2.0 * 6.0 + 3.0 * 7.0 + 4.0 * 8.0;

        assert!((result - expected).abs() < 1e-10);
    }

    #[test]
    fn test_simd_elementwise_add() {
        let simd_ops = SimdMatrixOps::new();
        let a = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let b = Array1::from_vec(vec![5.0, 6.0, 7.0, 8.0]);

        let result = simd_ops
            .elementwise_add_simd(&a, &b)
            .expect("operation should succeed");
        let expected = Array1::from_vec(vec![6.0, 8.0, 10.0, 12.0]);

        for (r, e) in result.iter().zip(expected.iter()) {
            assert!((r - e).abs() < 1e-10);
        }
    }

    #[test]
    fn test_simd_elementwise_mul() {
        let simd_ops = SimdMatrixOps::new();
        let a = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let b = Array1::from_vec(vec![5.0, 6.0, 7.0, 8.0]);

        let result = simd_ops
            .elementwise_mul_simd(&a, &b)
            .expect("operation should succeed");
        let expected = Array1::from_vec(vec![5.0, 12.0, 21.0, 32.0]);

        for (r, e) in result.iter().zip(expected.iter()) {
            assert!((r - e).abs() < 1e-10);
        }
    }

    #[test]
    fn test_matrix_vector_mul_simd() {
        let simd_ops = SimdMatrixOps::new();
        let matrix = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .expect("shape and data length should match");
        let vector = Array1::from_vec(vec![1.0, 2.0, 3.0]);

        let result = simd_ops
            .matrix_vector_mul_simd(&matrix, &vector)
            .expect("operation should succeed");
        let expected = Array1::from_vec(vec![14.0, 32.0]); // [1*1+2*2+3*3, 4*1+5*2+6*3]

        for (r, e) in result.iter().zip(expected.iter()) {
            assert!((r - e).abs() < 1e-10);
        }
    }

    #[test]
    fn test_parallel_operations() {
        let parallel_ops = ParallelDecomposition::new();

        // Test basic functionality
        let matrix = Array2::eye(3);
        let result = parallel_ops.parallel_eigendecomposition(&matrix);
        assert!(result.is_ok());

        let (eigenvals, eigenvecs) = result.expect("operation should succeed");
        assert_eq!(eigenvals.len(), 3);
        assert_eq!(eigenvecs.dim(), (3, 3));
    }

    #[test]
    fn test_mixed_precision() {
        let mixed_ops = MixedPrecisionOps::new();

        let input_f64 = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let converted_f32 = mixed_ops.to_single_precision(&input_f64);
        let converted_back = mixed_ops.to_double_precision(&converted_f32);

        for (orig, back) in input_f64.iter().zip(converted_back.iter()) {
            assert!((orig - back).abs() < 1e-6); // Precision loss expected
        }
    }

    #[test]
    fn test_aligned_memory() {
        let aligned_ops = AlignedMemoryOps::new(32);

        let aligned_vec = aligned_ops.create_aligned_array(10);
        assert_eq!(aligned_vec.len(), 10);
        assert!(aligned_ops.is_aligned(&aligned_vec));

        let array = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let aligned_array = aligned_ops.ensure_aligned(&array);
        assert_eq!(aligned_array.len(), array.len());
        assert!(aligned_ops.is_aligned(&aligned_array));

        assert_eq!(
            aligned_array.as_slice(),
            array.as_slice().expect("slice operation should succeed")
        );
    }

    #[test]
    fn test_vectorized_functions() {
        let simd_ops = SimdMatrixOps::new();
        let input = Array1::from_vec(vec![0.0, 1.0, 2.0]);

        let exp_result = simd_ops.vector_exp_simd(&input);
        let expected_exp = input.mapv(|x| x.exp());

        for (r, e) in exp_result.iter().zip(expected_exp.iter()) {
            assert!((r - e).abs() < 1e-10);
        }

        let sqrt_result = simd_ops.vector_sqrt_simd(&Array1::from_vec(vec![1.0, 4.0, 9.0]));
        let expected_sqrt = Array1::from_vec(vec![1.0, 2.0, 3.0]);

        for (r, e) in sqrt_result.iter().zip(expected_sqrt.iter()) {
            assert!((r - e).abs() < 1e-10);
        }
    }

    #[test]
    fn test_acceleration_config() {
        let config = AccelerationConfig {
            enable_simd: false,
            enable_parallel: false,
            enable_mixed_precision: true,
            enable_gpu: false,
            gpu_device_id: 0,
            gpu_memory_limit: None,
            num_threads: Some(4),
            memory_alignment: 64,
        };

        let simd_ops = SimdMatrixOps::new().with_config(config.clone());
        let a = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let b = Array1::from_vec(vec![4.0, 5.0, 6.0]);

        // Should use fallback when SIMD is disabled
        let result = simd_ops
            .dot_product_simd(&a, &b)
            .expect("operation should succeed");
        let expected = 32.0; // 1*4 + 2*5 + 3*6
        assert!((result - expected).abs() < 1e-10);
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_acceleration_creation() {
        // Test creation without GPU (should fallback gracefully)
        let _gpu_acc = GpuAcceleration::default();
        // This test just ensures the default creation works
        let _ = _gpu_acc;
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_config() {
        let config = AccelerationConfig {
            enable_gpu: true,
            gpu_device_id: 0,
            gpu_memory_limit: Some(1024 * 1024 * 1024), // 1GB
            ..AccelerationConfig::default()
        };

        // Test that config is properly set
        assert!(config.enable_gpu);
        assert_eq!(config.gpu_device_id, 0);
        assert_eq!(config.gpu_memory_limit, Some(1024 * 1024 * 1024));
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_matrix_multiply_identity() {
        let config = AccelerationConfig {
            enable_gpu: true,
            ..AccelerationConfig::default()
        };

        if let Ok(gpu_acc) = GpuAcceleration::with_config(config) {
            if gpu_acc.is_gpu_available() {
                let identity = Array2::from_shape_vec((2, 2), vec![1.0, 0.0, 0.0, 1.0])
                    .expect("shape and data length should match");
                let b = Array2::from_shape_vec((2, 2), vec![1.0, 2.0, 3.0, 4.0])
                    .expect("shape and data length should match");

                if let Ok(result) = gpu_acc.gpu_matrix_multiply(&identity, &b) {
                    assert_eq!(result.shape(), b.shape());
                    for (r, e) in result.iter().zip(b.iter()) {
                        assert!((r - e).abs() < 1e-10);
                    }
                }
            }
        }
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_decomposition_fallback() {
        // Test that GPU decomposition works with fallback
        let gpu_decomp = GpuDecomposition::default();

        let matrix =
            Array2::from_shape_vec((3, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
                .expect("operation should succeed");

        // This should work with fallback even without GPU
        if let Ok((factor_a, factor_b)) = gpu_decomp.gpu_factorize(&matrix) {
            assert_eq!(factor_a.nrows(), 3);
            assert_eq!(factor_b.ncols(), 3);
        }
    }

    #[cfg(feature = "gpu")]
    #[test]
    fn test_gpu_pca_basic() {
        let gpu_decomp = GpuDecomposition::default();

        let data = Array2::from_shape_vec(
            (4, 3),
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
            ],
        )
        .expect("operation should succeed");

        // Test GPU PCA with 2 components
        if let Ok((u, s, vt)) = gpu_decomp.gpu_pca(&data, 2) {
            assert_eq!(u.ncols(), 2);
            assert_eq!(s.len(), 2);
            assert_eq!(vt.nrows(), 2);
        }
    }
}
