//! CUDA kernel implementations for CFD operations
//!
//! This module provides real CUDA kernel compilation and execution using cudarc 0.17.3.
//! All kernels are compiled at runtime using NVRTC for maximum flexibility.

use super::*;
use cudarc::driver::{CudaContext, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::compile_ptx;
use std::collections::HashMap;
use std::sync::Arc;

/// CUDA kernel manager for CFD operations
/// Handles kernel compilation, memory management, and execution
pub struct CudaKernelManager {
    context: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    modules: HashMap<String, Arc<CudaModule>>,
}

impl CudaKernelManager {
    /// Create new CUDA kernel manager and compile all kernels
    pub fn new(config: &CfdConfig) -> CfdResult<Self> {
        let context = CudaContext::new(config.device_id as usize).map_err(|e| {
            CfdError::gpu_error(&format!("Failed to initialize CUDA context: {}", e))
        })?;
        let stream = context.default_stream();

        let mut manager = Self {
            context,
            stream,
            modules: HashMap::new(),
        };
        manager.compile_all_kernels()?;
        Ok(manager)
    }

    /// Compile all kernels at startup
    fn compile_all_kernels(&mut self) -> CfdResult<()> {
        // Compile advection kernels
        let advection_upwind_src = include_str!("cuda/advection_upwind.cu");
        self.compile_kernel_source("advection_kernels", advection_upwind_src)?;

        // Compile diffusion kernels
        let diffusion_src = include_str!("cuda/diffusion.cu");
        self.compile_kernel_source("diffusion_kernels", diffusion_src)?;

        // Compile 2D implicit diffusion kernels
        let diffusion_2d_src = include_str!("cuda/diffusion_2d_implicit.cu");
        self.compile_kernel_source("diffusion_2d_kernels", diffusion_2d_src)?;

        // Compile Poisson kernels
        let poisson_src = include_str!("cuda/poisson.cu");
        self.compile_kernel_source("poisson_kernels", poisson_src)?;

        // Compile matrix operation kernels
        let matrix_ops_src = include_str!("cuda/matrix_ops.cu");
        self.compile_kernel_source("matrix_kernels", matrix_ops_src)?;

        // Compile reduction kernels
        let reduction_src = include_str!("cuda/reduction.cu");
        self.compile_kernel_source("reduction_kernels", reduction_src)?;

        Ok(())
    }

    /// Compile a CUDA kernel from source code
    fn compile_kernel_source(&mut self, module_name: &str, source: &str) -> CfdResult<()> {
        // Compile to PTX using NVRTC
        let ptx = compile_ptx(source).map_err(|e| {
            CfdError::gpu_error(&format!("Failed to compile {} kernels: {}", module_name, e))
        })?;

        // Load PTX module into context
        let module = self.context.load_module(ptx).map_err(|e| {
            CfdError::gpu_error(&format!("Failed to load {} module: {}", module_name, e))
        })?;

        self.modules.insert(module_name.to_string(), module);
        Ok(())
    }

    /// Allocate GPU memory for f32 array
    pub fn allocate_f32(&self, size: usize) -> CfdResult<CudaSlice<f32>> {
        self.stream
            .alloc_zeros::<f32>(size)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))
    }

    /// Copy data from host to device
    pub fn copy_to_device(&self, host_data: &[f32]) -> CfdResult<CudaSlice<f32>> {
        self.stream
            .memcpy_stod(host_data)
            .map_err(|e| CfdError::gpu_error(&format!("Host to device copy failed: {}", e)))
    }

    /// Copy data from device to host
    pub fn copy_from_device(&self, device_slice: &CudaSlice<f32>) -> CfdResult<Vec<f32>> {
        self.stream
            .memcpy_dtov(device_slice)
            .map_err(|e| CfdError::gpu_error(&format!("Device to host copy failed: {}", e)))
    }

    /// Synchronize stream
    pub fn synchronize(&self) -> CfdResult<()> {
        self.stream
            .synchronize()
            .map_err(|e| CfdError::gpu_error(&format!("Stream synchronization failed: {}", e)))
    }

    /// Get context reference
    pub fn context(&self) -> &Arc<CudaContext> {
        &self.context
    }

    /// Get stream reference
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Get module by name
    pub fn get_module(&self, name: &str) -> CfdResult<&Arc<CudaModule>> {
        self.modules
            .get(name)
            .ok_or_else(|| CfdError::gpu_error(&format!("Module {} not found", name)))
    }

    /// Perform GPU reduction to compute sum
    pub fn reduce_sum(&self, input: &CudaSlice<f32>) -> CfdResult<f32> {
        let n = input.len();
        let block_size = 256;
        let grid_size = ((n as u32 + block_size - 1) / block_size).min(1024);

        // Allocate temporary buffer for partial results
        let mut partial_results = self.allocate_f32(grid_size as usize)?;

        // Load reduction module
        let module = self.get_module("reduction_kernels")?;

        // First pass: reduce to grid_size partial results
        let func = module
            .load_function("reduce_sum")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get reduce_sum kernel: {}", e)))?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(input)
                .arg(&mut partial_results)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Reduce sum kernel launch failed: {}", e))
                })?;
        }

        // Second pass: final reduction
        let final_func = module
            .load_function("reduce_final_sum")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get final sum kernel: {}", e)))?;

        let mut final_result = self.allocate_f32(1)?;

        let final_config = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
        };

        unsafe {
            self.stream
                .launch_builder(&final_func)
                .arg(&partial_results)
                .arg(&mut final_result)
                .arg(&(grid_size as i32))
                .launch(final_config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Final sum kernel launch failed: {}", e))
                })?;
        }

        self.synchronize()?;

        // Copy result back
        let mut result_host = vec![0.0f32; 1];
        self.stream
            .memcpy_dtoh(&final_result, &mut result_host)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy result from GPU: {}", e)))?;

        Ok(result_host[0])
    }

    /// Compute maximum absolute value using GPU reduction
    pub fn reduce_abs_max(&self, input: &CudaSlice<f32>) -> CfdResult<f32> {
        let n = input.len();
        let block_size = 256;
        let grid_size = ((n as u32 + block_size - 1) / block_size).min(1024);

        // Allocate temporary buffer for partial results
        let mut partial_results = self
            .stream
            .alloc_zeros::<f32>(grid_size as usize)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;

        let module = self.get_module("reduction_kernels")?;
        let reduce_func = module
            .load_function("reduce_abs_max")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get reduction kernel: {}", e)))?;

        // First reduction pass
        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.stream
                .launch_builder(&reduce_func)
                .arg(input)
                .arg(&mut partial_results)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Reduction kernel launch failed: {}", e))
                })?;
        }

        // Second reduction pass if needed
        if grid_size > 1 {
            let mut final_result = self
                .stream
                .alloc_zeros::<f32>(1)
                .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;

            let final_func = module.load_function("reduce_final_max").map_err(|e| {
                CfdError::gpu_error(&format!("Failed to get final reduction kernel: {}", e))
            })?;

            let final_config = LaunchConfig {
                grid_dim: (1, 1, 1),
                block_dim: (block_size, 1, 1),
                shared_mem_bytes: 0,
            };

            unsafe {
                self.stream
                    .launch_builder(&final_func)
                    .arg(&partial_results)
                    .arg(&mut final_result)
                    .arg(&(grid_size as i32))
                    .launch(final_config)
                    .map_err(|e| CfdError::gpu_error(&format!("Final reduction failed: {}", e)))?;
            }

            self.stream.synchronize()?;
            let result_host = self.copy_from_device(&final_result)?;
            Ok(result_host[0])
        } else {
            self.stream.synchronize()?;
            let result_host = self.copy_from_device(&partial_results)?;
            Ok(result_host[0])
        }
    }
}

/// Advection kernel wrapper for various schemes
pub struct AdvectionKernel {
    scheme: AdvectionScheme,
    stream: Arc<CudaStream>,
    module: Arc<CudaModule>,
}

impl AdvectionKernel {
    pub fn new(manager: &CudaKernelManager, scheme: AdvectionScheme) -> CfdResult<Self> {
        Ok(Self {
            scheme,
            stream: manager.stream.clone(),
            module: manager.get_module("advection_kernels")?.clone(),
        })
    }

    pub fn apply(
        &self,
        phi: &CudaSlice<f32>,
        phi_new: &mut CudaSlice<f32>,
        velocity: f32,
        dt: f32,
        dx: f32,
    ) -> CfdResult<()> {
        let n = phi.len();
        let grid_size = (n as u32 + 255) / 256;
        let block_size = 256;

        let kernel_name = match self.scheme {
            AdvectionScheme::Upwind => "advection_1d",
            AdvectionScheme::Central => "advection_central_1d",
            AdvectionScheme::Quick => "advection_quick_1d",
            AdvectionScheme::Weno => "advection_weno_1d",
        };

        let func = self.module.load_function(kernel_name).map_err(|e| {
            CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
        })?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(phi)
                .arg(phi_new)
                .arg(&velocity)
                .arg(&dt)
                .arg(&dx)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }

        self.stream
            .synchronize()
            .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
        Ok(())
    }

    pub fn apply_2d(
        &self,
        phi: &CudaSlice<f32>,
        phi_new: &mut CudaSlice<f32>,
        u: &CudaSlice<f32>,
        v: &CudaSlice<f32>,
        dt: f32,
        dx: f32,
        dy: f32,
        nx: usize,
        ny: usize,
    ) -> CfdResult<()> {
        let grid_dim_x = (nx as u32 + 15) / 16;
        let grid_dim_y = (ny as u32 + 15) / 16;

        let kernel_name = match self.scheme {
            AdvectionScheme::Upwind => "advection_2d",
            AdvectionScheme::Central => "advection_central_2d",
            AdvectionScheme::Quick => "advection_quick_2d",
            AdvectionScheme::Weno => "advection_weno_2d",
        };

        let func = self.module.load_function(kernel_name).map_err(|e| {
            CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
        })?;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, 1),
            block_dim: (16, 16, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(phi)
                .arg(phi_new)
                .arg(u)
                .arg(v)
                .arg(&dt)
                .arg(&dx)
                .arg(&dy)
                .arg(&(nx as i32))
                .arg(&(ny as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }

        self.stream
            .synchronize()
            .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
        Ok(())
    }
}

/// Diffusion kernel wrapper for various schemes
pub struct DiffusionKernel {
    scheme: DiffusionScheme,
    stream: Arc<CudaStream>,
    diffusion_module: Arc<CudaModule>,
    diffusion_2d_module: Arc<CudaModule>,
}

impl DiffusionKernel {
    pub fn new(manager: &CudaKernelManager, scheme: DiffusionScheme) -> CfdResult<Self> {
        Ok(Self {
            scheme,
            stream: manager.stream.clone(),
            diffusion_module: manager.get_module("diffusion_kernels")?.clone(),
            diffusion_2d_module: manager.get_module("diffusion_2d_kernels")?.clone(),
        })
    }

    pub fn apply(
        &self,
        temp: &CudaSlice<f32>,
        temp_new: &mut CudaSlice<f32>,
        alpha: f32,
        dt: f32,
        dx: f32,
    ) -> CfdResult<()> {
        let n = temp.len();
        let grid_size = (n as u32 + 255) / 256;
        let block_size = 256;

        let kernel_name = match self.scheme {
            DiffusionScheme::Explicit => "diffusion_explicit_1d",
            DiffusionScheme::Implicit => "diffusion_implicit_1d",
            DiffusionScheme::CrankNicolson => "diffusion_crank_nicolson_1d",
        };

        let func = self
            .diffusion_module
            .load_function(kernel_name)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
            })?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        match self.scheme {
            DiffusionScheme::Explicit => unsafe {
                self.stream
                    .launch_builder(&func)
                    .arg(temp)
                    .arg(temp_new)
                    .arg(&alpha)
                    .arg(&dt)
                    .arg(&dx)
                    .arg(&(n as i32))
                    .launch(config)
                    .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
            },
            DiffusionScheme::Implicit | DiffusionScheme::CrankNicolson => {
                // For implicit schemes, use iterative solver with convergence check
                let max_iterations = 100;
                let tolerance = 1e-6;
                let mut residual = 1.0;
                let mut iteration = 0;

                while residual > tolerance && iteration < max_iterations {
                    unsafe {
                        self.stream
                            .launch_builder(&func)
                            .arg(temp)
                            .arg(&mut *temp_new)
                            .arg(&alpha)
                            .arg(&dt)
                            .arg(&dx)
                            .arg(&(n as i32))
                            .launch(config)
                            .map_err(|e| {
                                CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
                            })?;
                    }
                    self.stream.synchronize().map_err(|e| {
                        CfdError::gpu_error(&format!("Synchronization failed: {}", e))
                    })?;

                    // TODO: Compute actual residual on GPU
                    residual *= 0.9; // Simplified convergence
                    iteration += 1;
                }
                return Ok(());
            }
        }

        self.stream
            .synchronize()
            .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
        Ok(())
    }

    pub fn apply_2d(
        &self,
        temp: &CudaSlice<f32>,
        temp_new: &mut CudaSlice<f32>,
        alpha: f32,
        dt: f32,
        dx: f32,
        dy: f32,
        nx: usize,
        ny: usize,
    ) -> CfdResult<()> {
        let grid_dim_x = (nx as u32 + 15) / 16;
        let grid_dim_y = (ny as u32 + 15) / 16;

        match self.scheme {
            DiffusionScheme::Explicit => {
                let func = self
                    .diffusion_module
                    .load_function("diffusion_explicit_2d")
                    .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;

                let config = LaunchConfig {
                    grid_dim: (grid_dim_x, grid_dim_y, 1),
                    block_dim: (16, 16, 1),
                    shared_mem_bytes: 0,
                };

                unsafe {
                    self.stream
                        .launch_builder(&func)
                        .arg(temp)
                        .arg(temp_new)
                        .arg(&alpha)
                        .arg(&dt)
                        .arg(&dx)
                        .arg(&dy)
                        .arg(&(nx as i32))
                        .arg(&(ny as i32))
                        .launch(config)
                        .map_err(|e| {
                            CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
                        })?;
                }

                self.stream
                    .synchronize()
                    .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
                Ok(())
            }
            DiffusionScheme::Implicit | DiffusionScheme::CrankNicolson => {
                // Implement 2D implicit/Crank-Nicolson using ADI (Alternating Direction Implicit)
                let kernel_name = match self.scheme {
                    DiffusionScheme::Implicit => "diffusion_implicit_2d_adi",
                    DiffusionScheme::CrankNicolson => "diffusion_crank_nicolson_2d_adi",
                    _ => unreachable!(),
                };

                let func = self
                    .diffusion_2d_module
                    .load_function(kernel_name)
                    .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;

                let config = LaunchConfig {
                    grid_dim: (grid_dim_x, grid_dim_y, 1),
                    block_dim: (16, 16, 1),
                    shared_mem_bytes: 0,
                };

                // ADI requires two sweeps: x-direction then y-direction
                unsafe {
                    self.stream
                        .launch_builder(&func)
                        .arg(temp)
                        .arg(temp_new)
                        .arg(&alpha)
                        .arg(&dt)
                        .arg(&dx)
                        .arg(&dy)
                        .arg(&(nx as i32))
                        .arg(&(ny as i32))
                        .launch(config)
                        .map_err(|e| {
                            CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
                        })?;
                }

                self.stream
                    .synchronize()
                    .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
                Ok(())
            }
        }
    }
}

/// Poisson equation solver kernel
pub struct PoissonKernel {
    stream: Arc<CudaStream>,
    module: Arc<CudaModule>,
    kernel_manager: Arc<CudaKernelManager>,
}

impl PoissonKernel {
    pub fn new(manager: &Arc<CudaKernelManager>) -> CfdResult<Self> {
        Ok(Self {
            stream: manager.stream.clone(),
            module: manager.get_module("poisson_kernels")?.clone(),
            kernel_manager: manager.clone(),
        })
    }

    /// [`Self::solve_2d`] under its older name.
    #[allow(clippy::too_many_arguments)]
    pub fn solve_jacobi_2d(
        &self,
        phi: &mut CudaSlice<f32>,
        source: &CudaSlice<f32>,
        nx: usize,
        ny: usize,
        dx: f32,
        dy: f32,
        max_iterations: usize,
        tolerance: f32,
    ) -> CfdResult<usize> {
        self.solve_2d(phi, source, nx, ny, dx, dy, max_iterations, tolerance)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn solve_2d(
        &self,
        phi: &mut CudaSlice<f32>,
        source: &CudaSlice<f32>,
        nx: usize,
        ny: usize,
        dx: f32,
        dy: f32,
        max_iterations: usize,
        tolerance: f32,
    ) -> CfdResult<usize> {
        let grid_dim_x = (nx as u32 + 15) / 16;
        let grid_dim_y = (ny as u32 + 15) / 16;

        // Calculate factors for Poisson equation
        let dx2_inv = 1.0f32 / (dx * dx);
        let dy2_inv = 1.0f32 / (dy * dy);
        let factor = 1.0f32 / (2.0f32 * (dx2_inv + dy2_inv));

        // Allocate temporary arrays for iteration
        let mut phi_temp = self
            .stream
            .alloc_zeros::<f32>(nx * ny)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
        let mut residual = self
            .stream
            .alloc_zeros::<f32>(nx * ny)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;

        let jacobi_func = self
            .module
            .load_function("poisson_jacobi_2d")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get Jacobi kernel: {}", e)))?;

        let residual_func = self
            .module
            .load_function("poisson_residual_2d")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get residual kernel: {}", e)))?;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, 1),
            block_dim: (16, 16, 1),
            shared_mem_bytes: 0,
        };

        for iteration in 0..max_iterations {
            // Perform Jacobi iteration
            unsafe {
                self.stream
                    .launch_builder(&jacobi_func)
                    .arg(&mut phi_temp)
                    .arg(&*phi)
                    .arg(source)
                    .arg(&factor)
                    .arg(&dx2_inv)
                    .arg(&dy2_inv)
                    .arg(&(nx as i32))
                    .arg(&(ny as i32))
                    .launch(config)
                    .map_err(|e| {
                        CfdError::gpu_error(&format!("Jacobi kernel launch failed: {}", e))
                    })?;
            }

            // Swap buffers
            std::mem::swap(&mut *phi, &mut phi_temp);

            // Check convergence every 10 iterations
            if iteration % 10 == 0 {
                unsafe {
                    self.stream
                        .launch_builder(&residual_func)
                        .arg(&mut residual)
                        .arg(&*phi)
                        .arg(source)
                        .arg(&dx2_inv)
                        .arg(&dy2_inv)
                        .arg(&(nx as i32))
                        .arg(&(ny as i32))
                        .launch(config)
                        .map_err(|e| {
                            CfdError::gpu_error(&format!("Residual kernel launch failed: {}", e))
                        })?;
                }

                // Use GPU reduction to find max residual
                let max_residual = self.kernel_manager.reduce_abs_max(&residual)?;

                if max_residual < tolerance {
                    return Ok(iteration + 1);
                }
            }

            self.stream
                .synchronize()
                .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
        }

        Ok(max_iterations)
    }
}

/// Matrix operations kernel for sparse linear algebra
pub struct MatrixOpsKernel {
    stream: Arc<CudaStream>,
    module: Arc<CudaModule>,
    kernel_manager: Arc<CudaKernelManager>,
}

impl MatrixOpsKernel {
    pub fn new(manager: &Arc<CudaKernelManager>) -> CfdResult<Self> {
        Ok(Self {
            stream: manager.stream.clone(),
            module: manager.get_module("matrix_kernels")?.clone(),
            kernel_manager: manager.clone(),
        })
    }

    pub fn dot_product(&self, x: &CudaSlice<f32>, y: &CudaSlice<f32>) -> CfdResult<f32> {
        let n = x.len();
        let grid_size = (n as u32 + 255) / 256;
        let block_size = 256;

        // Allocate result buffer for partial sums
        let mut partial_sums = self
            .stream
            .alloc_zeros::<f32>(grid_size as usize)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;

        let func = self
            .module
            .load_function("dot_product_partial")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: block_size as u32 * std::mem::size_of::<f32>() as u32,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(x)
                .arg(y)
                .arg(&mut partial_sums)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }

        // Use GPU reduction to compute final sum
        let result = self.kernel_manager.reduce_sum(&partial_sums)?;
        Ok(result)
    }

    /// `y = T x` for the tridiagonal matrix with `diagonal` (n) and
    /// `off_diagonal` (n − 1) — the `tridiagonal_matvec` kernel.
    pub fn tridiagonal_matvec(
        &self,
        diagonal: &CudaSlice<f32>,
        off_diagonal: &CudaSlice<f32>,
        x: &CudaSlice<f32>,
        y: &mut CudaSlice<f32>,
    ) -> CfdResult<()> {
        let n = x.len();
        let func = self
            .module
            .load_function("tridiagonal_matvec")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
        let config = LaunchConfig {
            grid_dim: ((n as u32 + 255) / 256, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(diagonal)
                .arg(off_diagonal)
                .arg(x)
                .arg(y)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }
        Ok(())
    }

    /// Compute L2 norm of vector
    pub fn vector_norm(&self, x: &CudaSlice<f32>) -> CfdResult<f32> {
        let n = x.len();
        let grid_size = (n as u32 + 255) / 256;
        let block_size = 256;

        let mut partial_sums = self
            .stream
            .alloc_zeros::<f32>(grid_size as usize)
            .map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;

        let func = self
            .module
            .load_function("l2_norm_partial")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: block_size as u32 * std::mem::size_of::<f32>() as u32,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(x)
                .arg(&mut partial_sums)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }

        let partial_host = self
            .stream
            .memcpy_dtov(&partial_sums)
            .map_err(|e| CfdError::gpu_error(&format!("Device to host copy failed: {}", e)))?;

        Ok(partial_host.iter().sum::<f32>().sqrt())
    }

    /// Compute y = alpha * x + y (AXPY operation)
    pub fn axpy(&self, alpha: f32, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>) -> CfdResult<()> {
        let n = x.len();
        let grid_size = (n as u32 + 255) / 256;
        let block_size = 256;

        let func = self
            .module
            .load_function("vector_axpy")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;

        let config = LaunchConfig {
            grid_dim: (grid_size, 1, 1),
            block_dim: (block_size, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.stream
                .launch_builder(&func)
                .arg(&alpha)
                .arg(x)
                .arg(y)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
        }

        self.stream
            .synchronize()
            .map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;

        Ok(())
    }
}