//! GPU-accelerated D2Q9 Lattice Boltzmann solver
//!
//! Implementation of the D2Q9 (2D, 9 velocities) LBM scheme on GPU
//! for incompressible flow simulations.

use crate::kernels::CudaKernelManager;
use crate::solvers::BoundaryConditions;
use crate::solvers::incompressible::flow_field::FlowField;
use crate::{CfdConfig, CfdError, CfdResult};
use cudarc::driver::{CudaModule, CudaSlice, LaunchConfig, PushKernelArg};
use nalgebra::Vector2;
use std::sync::Arc;

/// GPU buffers for D2Q9 LBM solver
struct D2Q9GpuBuffers {
    /// Distribution functions (9 velocities per cell)
    f: CudaSlice<f32>,
    /// Temporary distribution functions for streaming
    f_temp: CudaSlice<f32>,
    /// Equilibrium distributions
    f_eq: CudaSlice<f32>,
    /// Macroscopic density
    density: CudaSlice<f32>,
    /// Macroscopic velocity x-component
    velocity_x: CudaSlice<f32>,
    /// Macroscopic velocity y-component
    velocity_y: CudaSlice<f32>,
    /// Grid dimensions
    nx: usize,
    ny: usize,
}

/// GPU-accelerated D2Q9 Lattice Boltzmann solver
pub struct D2Q9GpuSolver {
    /// CUDA kernel manager
    kernel_manager: Arc<CudaKernelManager>,
    /// LBM module with compiled kernels
    lbm_module: Option<Arc<CudaModule>>,
    /// GPU buffers
    gpu_buffers: Option<D2Q9GpuBuffers>,
    /// Relaxation parameter
    omega: f64,
    /// Lattice speed
    cs2: f64,
}

impl D2Q9GpuSolver {
    /// Create a new D2Q9 GPU solver
    pub fn new(config: &CfdConfig) -> CfdResult<Self> {
        let kernel_manager = Arc::new(CudaKernelManager::new(config)?);

        Ok(Self {
            kernel_manager,
            lbm_module: None,
            gpu_buffers: None,
            omega: 1.0,     // Will be computed from viscosity
            cs2: 1.0 / 3.0, // Speed of sound squared for D2Q9
        })
    }

    /// Initialize solver with flow field
    pub fn initialize(&mut self, flow_field: &FlowField, viscosity: f64) -> CfdResult<()> {
        // Compute relaxation parameter from viscosity
        let dt = 1.0; // LBM time step
        let dx = 1.0; // LBM space step
        self.omega = 1.0 / (3.0 * viscosity * dt / (dx * dx) + 0.5);

        // Load LBM kernels
        self.load_lbm_kernels()?;

        // Initialize GPU buffers
        self.initialize_gpu_buffers(flow_field)?;

        Ok(())
    }

    /// Load and compile LBM kernels
    fn load_lbm_kernels(&mut self) -> CfdResult<()> {
        // For now, use a placeholder - in real implementation would load PTX
        // or compile CUDA kernels for D2Q9 operations

        // The module would contain kernels for:
        // - Collision step
        // - Streaming step
        // - Boundary conditions
        // - Equilibrium computation
        // - Macroscopic variables computation

        let module = self.kernel_manager.get_module("lbm_d2q9_kernels")?;
        self.lbm_module = Some(module.clone());

        Ok(())
    }

    /// Initialize GPU buffers from flow field
    fn initialize_gpu_buffers(&mut self, flow_field: &FlowField) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();
        let n_cells = nx * ny;
        let n_velocities = 9;
        let n_distributions = n_cells * n_velocities;

        // Allocate GPU memory
        let f = self.kernel_manager.allocate_f32(n_distributions)?;
        let f_temp = self.kernel_manager.allocate_f32(n_distributions)?;
        let f_eq = self.kernel_manager.allocate_f32(n_distributions)?;
        let density = self.kernel_manager.allocate_f32(n_cells)?;
        let velocity_x = self.kernel_manager.allocate_f32(n_cells)?;
        let velocity_y = self.kernel_manager.allocate_f32(n_cells)?;

        // Initialize distributions from flow field
        self.initialize_distributions_from_flow(
            &f,
            &density,
            &velocity_x,
            &velocity_y,
            flow_field,
        )?;

        self.gpu_buffers = Some(D2Q9GpuBuffers {
            f,
            f_temp,
            f_eq,
            density,
            velocity_x,
            velocity_y,
            nx,
            ny,
        });

        Ok(())
    }

    /// Initialize distributions from flow field data
    fn initialize_distributions_from_flow(
        &self,
        f: &CudaSlice<f32>,
        density: &CudaSlice<f32>,
        velocity_x: &CudaSlice<f32>,
        velocity_y: &CudaSlice<f32>,
        flow_field: &FlowField,
    ) -> CfdResult<()> {
        // Copy flow field data to GPU
        let (nx, ny, _, _) = flow_field.grid_info();

        // Flatten flow field data for GPU transfer
        let mut rho_host = Vec::with_capacity(nx * ny);
        let mut u_host = Vec::with_capacity(nx * ny);
        let mut v_host = Vec::with_capacity(nx * ny);

        for j in 0..ny {
            for i in 0..nx {
                // For incompressible flow, use constant density
                rho_host.push(1.0f32);
                // Get velocity from FlowField - note the staggered grid
                let (u_val, v_val) = flow_field.get_velocity_at(i, j).unwrap_or((0.0, 0.0));
                u_host.push(u_val as f32);
                v_host.push(v_val as f32);
            }
        }

        // Copy to GPU (htod: host source, device destination - needs mutable)
        let density_mut = &mut density.clone();
        let velocity_x_mut = &mut velocity_x.clone();
        let velocity_y_mut = &mut velocity_y.clone();

        self.kernel_manager
            .stream()
            .memcpy_htod(&rho_host, density_mut)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy density to GPU: {}", e)))?;
        self.kernel_manager
            .stream()
            .memcpy_htod(&u_host, velocity_x_mut)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_x to GPU: {}", e))
            })?;
        self.kernel_manager
            .stream()
            .memcpy_htod(&v_host, velocity_y_mut)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_y to GPU: {}", e))
            })?;

        // Initialize equilibrium distributions on GPU
        self.compute_equilibrium_distributions(f, density, velocity_x, velocity_y, nx, ny)?;

        Ok(())
    }

    /// Compute equilibrium distributions on GPU
    fn compute_equilibrium_distributions(
        &self,
        f: &CudaSlice<f32>,
        density: &CudaSlice<f32>,
        velocity_x: &CudaSlice<f32>,
        velocity_y: &CudaSlice<f32>,
        nx: usize,
        ny: usize,
    ) -> CfdResult<()> {
        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;

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

        let grid_dim_x = (nx as u32 + 15) / 16;
        let grid_dim_y = (ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut f.clone())
                .arg(density)
                .arg(velocity_x)
                .arg(velocity_y)
                .arg(&(self.cs2 as f32))
                .arg(&(nx as i32))
                .arg(&(ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Equilibrium init kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Execute one LBM time step
    pub async fn step(&mut self, boundary_conditions: &BoundaryConditions) -> CfdResult<()> {
        // Collision step
        self.gpu_collision_step()?;

        // Streaming step
        self.gpu_streaming_step()?;

        // Apply boundary conditions
        self.gpu_apply_boundaries(boundary_conditions)?;

        // Compute macroscopic variables
        self.gpu_compute_macroscopic()?;

        Ok(())
    }

    /// GPU-accelerated collision step
    fn gpu_collision_step(&self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        // First compute equilibrium distributions
        self.gpu_compute_equilibrium()?;

        // Then perform collision: f = f + omega * (f_eq - f)
        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module
            .load_function("d2q9_collision")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get collision kernel: {}", e)))?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.f.clone())
                .arg(&buffers.f_eq)
                .arg(&(self.omega as f32))
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Collision kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// GPU-accelerated streaming step
    pub fn gpu_streaming_step(&self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module
            .load_function("d2q9_streaming")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get streaming kernel: {}", e)))?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.f_temp.clone())
                .arg(&buffers.f)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Streaming kernel launch failed: {}", e))
                })?;
        }

        // Swap buffers: f = f_temp
        // In real implementation, would swap buffer pointers
        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// GPU-accelerated macroscopic variable computation
    fn gpu_compute_macroscopic(&self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module.load_function("d2q9_macroscopic").map_err(|e| {
            CfdError::gpu_error(&format!("Failed to get macroscopic kernel: {}", e))
        })?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.density.clone())
                .arg(&mut buffers.velocity_x.clone())
                .arg(&mut buffers.velocity_y.clone())
                .arg(&buffers.f)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!(
                        "Macroscopic variables kernel launch failed: {}",
                        e
                    ))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Compute equilibrium distributions on GPU
    fn gpu_compute_equilibrium(&self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module.load_function("d2q9_equilibrium").map_err(|e| {
            CfdError::gpu_error(&format!("Failed to get equilibrium kernel: {}", e))
        })?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.f_eq.clone())
                .arg(&buffers.density)
                .arg(&buffers.velocity_x)
                .arg(&buffers.velocity_y)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Equilibrium kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// GPU-accelerated bounce-back boundary conditions
    pub fn gpu_apply_bounce_back_boundaries(&self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module
            .load_function("d2q9_bounce_back_boundaries")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get boundary kernel: {}", e)))?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.f.clone())
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Boundary kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Complete GPU LBM time step
    pub fn gpu_step(&self) -> CfdResult<()> {
        self.gpu_collision_step()?;
        self.gpu_streaming_step()?;
        self.gpu_apply_bounce_back_boundaries()?;
        self.gpu_compute_macroscopic()?;
        Ok(())
    }

    /// Initialize flow with uniform velocity
    pub fn initialize_uniform_flow(
        &mut self,
        density: f64,
        velocity: Vector2<f64>,
    ) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .lbm_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("LBM module not loaded"))?;
        let func = module.load_function("d2q9_init_uniform").map_err(|e| {
            CfdError::gpu_error(&format!("Failed to get initialization kernel: {}", e))
        })?;

        let grid_dim_x = (buffers.nx as u32 + 15) / 16;
        let grid_dim_y = (buffers.ny as u32 + 15) / 16;

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

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.f.clone())
                .arg(&mut buffers.density.clone())
                .arg(&mut buffers.velocity_x.clone())
                .arg(&mut buffers.velocity_y.clone())
                .arg(&(density as f32))
                .arg(&(velocity.x as f32))
                .arg(&(velocity.y as f32))
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Initialization kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Get macroscopic variables at a specific point (copy from GPU)
    pub fn get_macroscopic_at(&self, i: usize, j: usize) -> CfdResult<(f64, Vector2<f64>)> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let idx = j * buffers.nx + i;

        // Copy single values from GPU
        let mut density_val = vec![0.0f32; 1];
        let mut vx_val = vec![0.0f32; 1];
        let mut vy_val = vec![0.0f32; 1];

        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.density.slice(idx..idx + 1), &mut density_val)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy density from GPU: {}", e)))?;
        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.velocity_x.slice(idx..idx + 1), &mut vx_val)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_x from GPU: {}", e))
            })?;
        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.velocity_y.slice(idx..idx + 1), &mut vy_val)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_y from GPU: {}", e))
            })?;

        self.kernel_manager.synchronize()?;

        Ok((
            density_val[0] as f64,
            Vector2::new(vx_val[0] as f64, vy_val[0] as f64),
        ))
    }

    /// Apply boundary conditions
    fn gpu_apply_boundaries(&self, _boundary_conditions: &BoundaryConditions) -> CfdResult<()> {
        // For now, just apply bounce-back
        self.gpu_apply_bounce_back_boundaries()?;
        Ok(())
    }

    /// Update flow field from GPU results
    pub fn update_flow_field(&self, flow_field: &mut FlowField) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let (nx, ny, _, _) = flow_field.grid_info();
        let n_cells = nx * ny;

        // Copy results back from GPU
        let mut density_host = vec![0.0f32; n_cells];
        let mut vx_host = vec![0.0f32; n_cells];
        let mut vy_host = vec![0.0f32; n_cells];

        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.density, &mut density_host)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy density from GPU: {}", e)))?;
        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.velocity_x, &mut vx_host)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_x from GPU: {}", e))
            })?;
        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.velocity_y, &mut vy_host)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy velocity_y from GPU: {}", e))
            })?;

        self.kernel_manager.synchronize()?;

        // Update flow field velocities
        for j in 0..ny {
            for i in 0..nx {
                let idx = j * nx + i;
                // Note: density is not stored in incompressible FlowField
                // Update velocities using set_velocity method
                flow_field.set_velocity(i, j, vx_host[idx] as f64, vy_host[idx] as f64)?;
            }
        }

        Ok(())
    }
}