//! Fused GPU reduction kernels: stats, argmax, sum.
//!
//! Computes multiple statistics in a single kernel launch with shared-memory
//! reductions, eliminating per-stat GPU->CPU readbacks. All computation is
//! GPU-resident with a single DtoH transfer for the final scalar results.
//!
//! ## Kernels
//!
//! - `fused_stats_reduce`: min, max, sum, sum_sq, count in one pass
//! - `argmax_flat`: argmax of a flat buffer -> single u32 index
//! - `argmax_rows`: argmax per row for 2D data [rows, cols] -> u32 per row
//! - `sum_reduce`: sum of all elements -> single scalar
//! - `col_sum_reduce`: per-column sum for 2D data [rows, cols] -> f32 per col

#![allow(unsafe_code)] // CUDA FFI requires unsafe for kernel launches.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use crate::MLError;

/// Compiled reduction kernels — cached per CUDA context.
#[allow(missing_debug_implementations)]
pub struct ReductionKernels {
    fused_stats_fn: CudaFunction,
    argmax_flat_fn: CudaFunction,
    argmax_rows_fn: CudaFunction,
    sum_reduce_fn: CudaFunction,
    col_sum_fn: CudaFunction,
    stream: Arc<CudaStream>,
}

/// Result of a fused statistics computation.
#[derive(Debug, Clone, Copy)]
pub struct Stats {
    /// Minimum element value.
    pub min: f32,
    /// Maximum element value.
    pub max: f32,
    /// Arithmetic mean.
    pub mean: f32,
    /// Population variance.
    pub variance: f32,
    /// Number of elements.
    pub count: usize,
}

impl ReductionKernels {
    /// Compile all reduction kernels.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(REDUCTION_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("reduction kernel compilation: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("reduction module load: {e}"))
        })?;

        let load = |name: &str| -> Result<CudaFunction, MLError> {
            module.load_function(name).map_err(|e| {
                MLError::ModelError(format!("reduction kernel '{name}' load: {e}"))
            })
        };

        Ok(Self {
            fused_stats_fn: load("fused_stats_reduce")?,
            argmax_flat_fn: load("argmax_flat")?,
            argmax_rows_fn: load("argmax_rows")?,
            sum_reduce_fn: load("sum_reduce")?,
            col_sum_fn: load("col_sum_reduce")?,
            stream: Arc::clone(stream),
        })
    }

    /// Compute min, max, mean, variance in a single kernel launch + single DtoH readback.
    ///
    /// The kernel performs a tree-reduction in shared memory (warp-level then block-level)
    /// with atomic writes to a 5-element global result buffer: [min, max, sum, sum_sq, count].
    /// Mean and variance are derived on the CPU from the 5 raw values (20 bytes transferred).
    pub fn stats(&self, data: &CudaSlice<f32>, n: usize) -> Result<Stats, MLError> {
        if n == 0 {
            return Err(MLError::InvalidInput("stats: empty buffer".to_string()));
        }

        // result[5] = [min, max, sum, sum_sq, count]
        let mut result = self.stream.alloc_zeros::<f32>(5).map_err(|e| {
            MLError::ModelError(format!("stats result alloc: {e}"))
        })?;

        // Initialize min to +INF, max to -INF, sum/sum_sq/count to 0
        let init_vals = [f32::INFINITY, f32::NEG_INFINITY, 0.0_f32, 0.0_f32, 0.0_f32];
        self.stream.memcpy_htod(&init_vals, &mut result).map_err(|e| {
            MLError::ModelError(format!("stats init upload: {e}"))
        })?;

        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        // Cap blocks to avoid excessive atomic contention
        let blocks = blocks.min(1024);
        let n_i32 = n as i32;
        let shared_bytes = threads * 5 * 4; // 5 floats per thread for shared reduction (min, max, sum, sum_sq, count)

        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: shared_bytes,
        };

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* data, float* result, int n)
        // data is a valid CudaSlice<f32> of at least n elements.
        // result is a valid CudaSlice<f32> of 5 elements, pre-initialized.
        unsafe {
            self.stream
                .launch_builder(&self.fused_stats_fn)
                .arg(data)
                .arg(&result)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("fused_stats_reduce: {e}")))?;
        }

        // Single DtoH readback: 5 floats = 20 bytes
        let mut host = [0.0_f32; 5];
        self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
            MLError::ModelError(format!("stats DtoH readback: {e}"))
        })?;

        let min_val = host.first().copied().unwrap_or(0.0);
        let max_val = host.get(1).copied().unwrap_or(0.0);
        let sum_val = host.get(2).copied().unwrap_or(0.0);
        let sum_sq_val = host.get(3).copied().unwrap_or(0.0);
        let count_val = host.get(4).copied().unwrap_or(0.0);

        let count_f = if count_val > 0.0 { count_val } else { 1.0 };
        let mean = sum_val / count_f;
        let variance = (sum_sq_val / count_f) - (mean * mean);
        // Clamp variance to 0 in case of floating point rounding
        let variance = if variance < 0.0 { 0.0 } else { variance };

        Ok(Stats {
            min: min_val,
            max: max_val,
            mean,
            variance,
            count: count_val as usize,
        })
    }

    /// Argmax of a flat buffer -- returns index as u32.
    ///
    /// Single kernel launch with shared-memory reduction, single 4-byte DtoH readback.
    pub fn argmax(&self, data: &CudaSlice<f32>, n: usize) -> Result<u32, MLError> {
        if n == 0 {
            return Err(MLError::InvalidInput("argmax: empty buffer".to_string()));
        }

        // Output: [best_val, best_idx] as float2 (index packed into float bits)
        let mut result = self.stream.alloc_zeros::<f32>(2).map_err(|e| {
            MLError::ModelError(format!("argmax result alloc: {e}"))
        })?;
        let init_vals = [f32::NEG_INFINITY, 0.0_f32]; // val, idx
        self.stream.memcpy_htod(&init_vals, &mut result).map_err(|e| {
            MLError::ModelError(format!("argmax init upload: {e}"))
        })?;

        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        let blocks = blocks.min(1024);
        let n_i32 = n as i32;
        let shared_bytes = threads * 2 * 4; // 2 floats per thread (val + idx)

        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: shared_bytes,
        };

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* data, float* result, int n)
        // data is a valid CudaSlice<f32> of at least n elements.
        // result is a valid CudaSlice<f32> of 2 elements, pre-initialized.
        unsafe {
            self.stream
                .launch_builder(&self.argmax_flat_fn)
                .arg(data)
                .arg(&result)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("argmax_flat: {e}")))?;
        }

        let mut host = [0.0_f32; 2];
        self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
            MLError::ModelError(format!("argmax DtoH readback: {e}"))
        })?;

        // Index is stored as __int_as_float in the kernel
        Ok(f32::to_bits(host.get(1).copied().unwrap_or(0.0)))
    }

    /// Argmax per row for 2D data [rows, cols] -- returns `Vec<u32>` of length `rows`.
    ///
    /// Each block processes one row. Single DtoH readback of `rows * 4` bytes.
    pub fn argmax_rows(&self, data: &CudaSlice<f32>, rows: usize, cols: usize) -> Result<Vec<u32>, MLError> {
        if rows == 0 || cols == 0 {
            return Err(MLError::InvalidInput("argmax_rows: empty data".to_string()));
        }

        let result = self.stream.alloc_zeros::<u32>(rows).map_err(|e| {
            MLError::ModelError(format!("argmax_rows result alloc: {e}"))
        })?;

        let threads = 256_u32.min(cols as u32);
        let rows_i32 = rows as i32;
        let cols_i32 = cols as i32;
        let shared_bytes = threads * 2 * 4; // 2 floats per thread (val + idx)

        let cfg = LaunchConfig {
            grid_dim: (rows as u32, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: shared_bytes,
        };

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* data, unsigned int* result, int rows, int cols)
        // data is a valid CudaSlice<f32> of at least rows*cols elements.
        // result is a valid CudaSlice<u32> of at least rows elements.
        unsafe {
            self.stream
                .launch_builder(&self.argmax_rows_fn)
                .arg(data)
                .arg(&result)
                .arg(&rows_i32)
                .arg(&cols_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("argmax_rows: {e}")))?;
        }

        let mut host = vec![0_u32; rows];
        self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
            MLError::ModelError(format!("argmax_rows DtoH readback: {e}"))
        })?;

        Ok(host)
    }

    /// Sum of all elements -- single scalar readback.
    ///
    /// Tree-reduction in shared memory with atomic add to global accumulator.
    pub fn sum(&self, data: &CudaSlice<f32>, n: usize) -> Result<f32, MLError> {
        if n == 0 {
            return Err(MLError::InvalidInput("sum: empty buffer".to_string()));
        }

        let mut result = self.stream.alloc_zeros::<f32>(1).map_err(|e| {
            MLError::ModelError(format!("sum result alloc: {e}"))
        })?;
        self.stream.memset_zeros(&mut result).map_err(|e| {
            MLError::ModelError(format!("sum result zero: {e}"))
        })?;

        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        let blocks = blocks.min(1024);
        let n_i32 = n as i32;
        let shared_bytes = threads * 4; // 1 float per thread

        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: shared_bytes,
        };

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* data, float* result, int n)
        // data is a valid CudaSlice<f32> of at least n elements.
        // result is a valid CudaSlice<f32> of 1 element, zero-initialized.
        unsafe {
            self.stream
                .launch_builder(&self.sum_reduce_fn)
                .arg(data)
                .arg(&result)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("sum_reduce: {e}")))?;
        }

        let mut host = [0.0_f32];
        self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
            MLError::ModelError(format!("sum DtoH readback: {e}"))
        })?;

        Ok(host.first().copied().unwrap_or(0.0))
    }

    /// Per-column sum for 2D data [rows, cols] -- returns `Vec<f32>` of length `cols`.
    ///
    /// One block per column, shared-memory tree reduction along the row axis.
    pub fn col_sums(&self, data: &CudaSlice<f32>, rows: usize, cols: usize) -> Result<Vec<f32>, MLError> {
        if rows == 0 || cols == 0 {
            return Err(MLError::InvalidInput("col_sums: empty data".to_string()));
        }

        let mut result = self.stream.alloc_zeros::<f32>(cols).map_err(|e| {
            MLError::ModelError(format!("col_sums result alloc: {e}"))
        })?;
        self.stream.memset_zeros(&mut result).map_err(|e| {
            MLError::ModelError(format!("col_sums result zero: {e}"))
        })?;

        let threads = 256_u32.min(rows as u32);
        let rows_i32 = rows as i32;
        let cols_i32 = cols as i32;
        let shared_bytes = threads * 4; // 1 float per thread

        let cfg = LaunchConfig {
            grid_dim: (cols as u32, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: shared_bytes,
        };

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* data, float* result, int rows, int cols)
        // data is a valid CudaSlice<f32> of at least rows*cols elements.
        // result is a valid CudaSlice<f32> of at least cols elements.
        unsafe {
            self.stream
                .launch_builder(&self.col_sum_fn)
                .arg(data)
                .arg(&result)
                .arg(&rows_i32)
                .arg(&cols_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("col_sum_reduce: {e}")))?;
        }

        let mut host = vec![0.0_f32; cols];
        self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
            MLError::ModelError(format!("col_sums DtoH readback: {e}"))
        })?;

        Ok(host)
    }
}

// ── CUDA source ──────────────────────────────────────────────────────────

const REDUCTION_CUDA_SRC: &str = r#"
// ═══════════════════════════════════════════════════════════════════════
// Fused reduction kernels: stats, argmax, sum, per-column sum.
//
// All kernels use shared-memory tree reduction (standard CUDA pattern):
// 1. Each thread loads and processes its elements
// 2. Write thread-local results to shared memory
// 3. Tree-reduce within the block (stride halving)
// 4. Block-level result atomically merged into global output
//
// Warp-level primitives (__shfl_down_sync) are used for the final
// warp to avoid shared-memory bank conflicts.
// ═══════════════════════════════════════════════════════════════════════

// ── Helper: warp-level reductions ─────────────────────────────────────

__device__ __forceinline__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xFFFFFFFF, val, offset);
    return val;
}

__device__ __forceinline__ float warp_reduce_min(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fminf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
    return val;
}

__device__ __forceinline__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
    return val;
}

// ── Fused stats: min, max, sum, sum_sq, count ─────────────────────────

extern "C" __global__
void fused_stats_reduce(const float* __restrict__ data,
                        float* __restrict__ result,
                        int n) {
    // Shared memory layout: [blockDim.x * 4]
    // offset 0: local_min, offset B: local_max, offset 2B: local_sum, offset 3B: local_sum_sq
    extern __shared__ float shmem[];
    int B = blockDim.x;

    float local_min = 3.402823466e+38f;  // FLT_MAX
    float local_max = -3.402823466e+38f; // -FLT_MAX
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    int local_count = 0;

    // Grid-stride loop: each thread processes multiple elements
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        float v = data[i];
        local_min = fminf(local_min, v);
        local_max = fmaxf(local_max, v);
        local_sum += v;
        local_sum_sq += v * v;
        local_count++;
    }

    // Store in shared memory
    shmem[threadIdx.x]         = local_min;
    shmem[B + threadIdx.x]     = local_max;
    shmem[2*B + threadIdx.x]   = local_sum;
    shmem[3*B + threadIdx.x]   = local_sum_sq;
    shmem[4*B + threadIdx.x]   = (float)local_count;
    __syncthreads();

    // Tree reduction in shared memory
    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem[threadIdx.x]         = fminf(shmem[threadIdx.x], shmem[threadIdx.x + s]);
            shmem[B + threadIdx.x]     = fmaxf(shmem[B + threadIdx.x], shmem[B + threadIdx.x + s]);
            shmem[2*B + threadIdx.x]  += shmem[2*B + threadIdx.x + s];
            shmem[3*B + threadIdx.x]  += shmem[3*B + threadIdx.x + s];
            shmem[4*B + threadIdx.x]  += shmem[4*B + threadIdx.x + s];
        }
        __syncthreads();
    }

    // Final warp reduction (no __syncthreads needed within a warp)
    if (threadIdx.x < 32) {
        // Fold the second set of 32 elements (indices 32..63) into 0..31
        float w_min = fminf(shmem[threadIdx.x], shmem[threadIdx.x + 32]);
        float w_max = fmaxf(shmem[B + threadIdx.x], shmem[B + threadIdx.x + 32]);
        float w_sum = shmem[2*B + threadIdx.x] + shmem[2*B + threadIdx.x + 32];
        float w_ssq = shmem[3*B + threadIdx.x] + shmem[3*B + threadIdx.x + 32];
        float w_cnt = shmem[4*B + threadIdx.x] + shmem[4*B + threadIdx.x + 32];

        // Reduce within the warp using shuffle
        w_min = warp_reduce_min(w_min);
        w_max = warp_reduce_max(w_max);
        w_sum = warp_reduce_sum(w_sum);
        w_ssq = warp_reduce_sum(w_ssq);
        w_cnt = warp_reduce_sum(w_cnt);

        if (threadIdx.x == 0) {
            // Atomic merge to global result
            // For simplicity and correctness with negative values, use CAS loops.
            {
                // atomicMin for float via CAS loop
                unsigned int* addr = (unsigned int*)&result[0];
                unsigned int old_val = *addr;
                unsigned int new_val = __float_as_uint(w_min);
                while (__uint_as_float(old_val) > w_min) {
                    unsigned int assumed = old_val;
                    old_val = atomicCAS(addr, assumed, new_val);
                    if (old_val == assumed) break;
                    new_val = __float_as_uint(fminf(__uint_as_float(old_val), w_min));
                }
            }
            {
                // atomicMax for float via CAS loop
                unsigned int* addr = (unsigned int*)&result[1];
                unsigned int old_val = *addr;
                unsigned int new_val = __float_as_uint(w_max);
                while (__uint_as_float(old_val) < w_max) {
                    unsigned int assumed = old_val;
                    old_val = atomicCAS(addr, assumed, new_val);
                    if (old_val == assumed) break;
                    new_val = __float_as_uint(fmaxf(__uint_as_float(old_val), w_max));
                }
            }
            atomicAdd(&result[2], w_sum);
            atomicAdd(&result[3], w_ssq);
            atomicAdd(&result[4], w_cnt);
        }
    }
}

// ── Argmax (flat buffer) ──────────────────────────────────────────────

extern "C" __global__
void argmax_flat(const float* __restrict__ data,
                 float* __restrict__ result,
                 int n) {
    // Shared memory: [blockDim.x * 2] -- (val, idx_as_float) pairs
    extern __shared__ float shmem[];
    int B = blockDim.x;

    float local_val = -3.402823466e+38f;
    int local_idx = 0;

    // Grid-stride loop
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        float v = data[i];
        if (v > local_val) {
            local_val = v;
            local_idx = i;
        }
    }

    shmem[threadIdx.x] = local_val;
    shmem[B + threadIdx.x] = __int_as_float(local_idx);
    __syncthreads();

    // Tree reduction
    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            if (shmem[threadIdx.x + s] > shmem[threadIdx.x]) {
                shmem[threadIdx.x] = shmem[threadIdx.x + s];
                shmem[B + threadIdx.x] = shmem[B + threadIdx.x + s];
            }
        }
        __syncthreads();
    }

    // Final warp
    if (threadIdx.x < 32) {
        // Fold shmem[32..63] into shmem[0..31]
        float w_val = shmem[threadIdx.x];
        float w_idx = shmem[B + threadIdx.x];
        float other_val32 = shmem[threadIdx.x + 32];
        float other_idx32 = shmem[B + threadIdx.x + 32];
        if (other_val32 > w_val) {
            w_val = other_val32;
            w_idx = other_idx32;
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
            float other_idx = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
            if (other_val > w_val) {
                w_val = other_val;
                w_idx = other_idx;
            }
        }

        if (threadIdx.x == 0) {
            // CAS loop to merge block results into global output
            unsigned int* val_addr = (unsigned int*)&result[0];
            unsigned int old_val_bits = *val_addr;
            while (__uint_as_float(old_val_bits) < w_val) {
                unsigned int assumed = old_val_bits;
                // Pack both val and idx atomically
                old_val_bits = atomicCAS(val_addr, assumed, __float_as_uint(w_val));
                if (old_val_bits == assumed) {
                    // We won the race -- write the index
                    result[1] = w_idx;
                    break;
                }
            }
        }
    }
}

// ── Argmax per row for 2D data ────────────────────────────────────────

extern "C" __global__
void argmax_rows(const float* __restrict__ data,
                 unsigned int* __restrict__ result,
                 int rows,
                 int cols) {
    // One block per row: blockIdx.x = row index
    extern __shared__ float shmem[];
    int B = blockDim.x;
    int row = blockIdx.x;
    if (row >= rows) return;

    const float* row_data = data + row * cols;

    float local_val = -3.402823466e+38f;
    int local_idx = 0;

    // Stride loop within the row
    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
        float v = row_data[c];
        if (v > local_val) {
            local_val = v;
            local_idx = c;
        }
    }

    shmem[threadIdx.x] = local_val;
    shmem[B + threadIdx.x] = __int_as_float(local_idx);
    __syncthreads();

    // Tree reduction within block
    for (int s = B >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            if (shmem[threadIdx.x + s] > shmem[threadIdx.x]) {
                shmem[threadIdx.x] = shmem[threadIdx.x + s];
                shmem[B + threadIdx.x] = shmem[B + threadIdx.x + s];
            }
        }
        __syncthreads();
    }

    // Final warp
    if (threadIdx.x < 32) {
        float w_val = shmem[threadIdx.x];
        float w_idx = shmem[B + threadIdx.x];
        // Fold shmem[32..63] into shmem[0..31] when block is large enough
        if (B >= 64) {
            float other_val32 = shmem[threadIdx.x + 32];
            float other_idx32 = shmem[B + threadIdx.x + 32];
            if (other_val32 > w_val) {
                w_val = other_val32;
                w_idx = other_idx32;
            }
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
            float other_idx = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
            if (other_val > w_val) {
                w_val = other_val;
                w_idx = other_idx;
            }
        }

        if (threadIdx.x == 0) {
            result[row] = (unsigned int)__float_as_int(w_idx);
        }
    }
}

// ── Sum reduction ─────────────────────────────────────────────────────

extern "C" __global__
void sum_reduce(const float* __restrict__ data,
                float* __restrict__ result,
                int n) {
    extern __shared__ float shmem[];

    float local_sum = 0.0f;

    // Grid-stride loop
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
        local_sum += data[i];
    }

    shmem[threadIdx.x] = local_sum;
    __syncthreads();

    // Tree reduction
    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem[threadIdx.x] += shmem[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        // Fold shmem[32..63] into shmem[0..31]
        float w_sum = shmem[threadIdx.x] + shmem[threadIdx.x + 32];
        w_sum = warp_reduce_sum(w_sum);
        if (threadIdx.x == 0) {
            atomicAdd(result, w_sum);
        }
    }
}

// ── Per-column sum for 2D data ────────────────────────────────────────

extern "C" __global__
void col_sum_reduce(const float* __restrict__ data,
                    float* __restrict__ result,
                    int rows,
                    int cols) {
    // One block per column: blockIdx.x = column index
    extern __shared__ float shmem[];
    int col = blockIdx.x;
    if (col >= cols) return;

    float local_sum = 0.0f;

    // Stride loop along the column (row axis)
    for (int r = threadIdx.x; r < rows; r += blockDim.x) {
        local_sum += data[r * cols + col];
    }

    shmem[threadIdx.x] = local_sum;
    __syncthreads();

    // Tree reduction
    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            shmem[threadIdx.x] += shmem[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        // Fold shmem[32..63] into shmem[0..31] when block is large enough
        float w_sum = shmem[threadIdx.x];
        if (blockDim.x >= 64)
            w_sum += shmem[threadIdx.x + 32];
        w_sum = warp_reduce_sum(w_sum);
        if (threadIdx.x == 0) {
            result[col] = w_sum;
        }
    }
}
"#;

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
            .expect("CUDA required")
            .new_stream()
            .expect("fork stream")
    }

    #[test]
    fn test_stats_basic() {
        let stream = make_stream();
        let kernels = ReductionKernels::new(&stream).unwrap();

        let data: Vec<f32> = (1..=100).map(|i| i as f32).collect();
        let mut gpu_data = stream.alloc_zeros::<f32>(100).unwrap();
        stream.memcpy_htod(&data, &mut gpu_data).unwrap();

        let stats = kernels.stats(&gpu_data, 100).unwrap();
        assert!((stats.min - 1.0).abs() < 1e-3, "min={}", stats.min);
        assert!((stats.max - 100.0).abs() < 1e-3, "max={}", stats.max);
        assert!((stats.mean - 50.5).abs() < 0.1, "mean={}", stats.mean);
        assert_eq!(stats.count, 100);
        // Variance of 1..100 = (100^2 - 1) / 12 = 833.25
        assert!((stats.variance - 833.25).abs() < 2.0, "var={}", stats.variance);
    }

    #[test]
    fn test_argmax_flat() {
        let stream = make_stream();
        let kernels = ReductionKernels::new(&stream).unwrap();

        let mut data = vec![0.0_f32; 1000];
        // Place max at index 42
        if let Some(v) = data.get_mut(42) {
            *v = 999.0;
        }
        let mut gpu_data = stream.alloc_zeros::<f32>(1000).unwrap();
        stream.memcpy_htod(&data, &mut gpu_data).unwrap();

        let idx = kernels.argmax(&gpu_data, 1000).unwrap();
        assert_eq!(idx, 42, "argmax should be 42, got {idx}");
    }

    #[test]
    fn test_argmax_rows() {
        let stream = make_stream();
        let kernels = ReductionKernels::new(&stream).unwrap();

        // 4 rows x 5 cols, max at different columns per row
        #[allow(clippy::approx_constant)]
        let data: Vec<f32> = vec![
            1.0, 5.0, 2.0, 3.0, 4.0,  // max at col 1
            9.0, 1.0, 2.0, 3.0, 4.0,  // max at col 0
            1.0, 2.0, 3.0, 4.0, 8.0,  // max at col 4
            1.0, 2.0, 7.0, 3.0, 4.0,  // max at col 2
        ];
        let mut gpu_data = stream.alloc_zeros::<f32>(20).unwrap();
        stream.memcpy_htod(&data, &mut gpu_data).unwrap();

        let indices = kernels.argmax_rows(&gpu_data, 4, 5).unwrap();
        assert_eq!(indices, vec![1, 0, 4, 2], "argmax_rows={indices:?}");
    }

    #[test]
    fn test_sum_basic() {
        let stream = make_stream();
        let kernels = ReductionKernels::new(&stream).unwrap();

        let data: Vec<f32> = (1..=100).map(|i| i as f32).collect();
        let mut gpu_data = stream.alloc_zeros::<f32>(100).unwrap();
        stream.memcpy_htod(&data, &mut gpu_data).unwrap();

        let total = kernels.sum(&gpu_data, 100).unwrap();
        assert!((total - 5050.0).abs() < 1.0, "sum={total}");
    }

    #[test]
    fn test_col_sums() {
        let stream = make_stream();
        let kernels = ReductionKernels::new(&stream).unwrap();

        // 3 rows x 4 cols
        let data: Vec<f32> = 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,
        ];
        let mut gpu_data = stream.alloc_zeros::<f32>(12).unwrap();
        stream.memcpy_htod(&data, &mut gpu_data).unwrap();

        let sums = kernels.col_sums(&gpu_data, 3, 4).unwrap();
        // col 0: 1+5+9=15, col 1: 2+6+10=18, col 2: 3+7+11=21, col 3: 4+8+12=24
        assert_eq!(sums.len(), 4);
        assert!((sums.first().copied().unwrap_or(0.0) - 15.0).abs() < 1e-3, "col0={:?}", sums.first());
        assert!((sums.get(1).copied().unwrap_or(0.0) - 18.0).abs() < 1e-3);
        assert!((sums.get(2).copied().unwrap_or(0.0) - 21.0).abs() < 1e-3);
        assert!((sums.get(3).copied().unwrap_or(0.0) - 24.0).abs() < 1e-3);
    }
}
