//! Element-wise and memory-layout CUDA kernels for `GpuTensor` operations.
//!
//! Replaces host-roundtrip workarounds with GPU-native kernels:
//!
//! - **Binary ops**: add, sub, mul (parameterized via `op` arg)
//! - **Unary ops**: powf, sqr, floor, relu, clamp (parameterized via `op` arg)
//! - **Broadcast ops**: scalar/row broadcast mul, div, expand
//! - **Memory ops**: transpose, narrow (dim-0 DtoD), index_select (gather)
//!
//! All kernels are compiled once via `compile_ptx_for_device()` and cached in
//! a `OnceLock`. Thread-safe, no global mutable state.

#![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;

/// Launch config helper for element-wise kernels.
fn elem_cfg(n: usize) -> LaunchConfig {
    let threads = 256_u32;
    let blocks = ((n as u32) + threads - 1) / threads;
    LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (threads, 1, 1),
        shared_mem_bytes: 0,
    }
}

/// Compiled element-wise kernels -- cached per CUDA context.
///
/// Construct once via [`ElementwiseKernels::new`], then reuse for all
/// element-wise operations on tensors sharing the same context.
#[allow(missing_debug_implementations)]
pub struct ElementwiseKernels {
    binary_fn: CudaFunction,
    unary_fn: CudaFunction,
    broadcast_binary_fn: CudaFunction,
    broadcast_row_binary_fn: CudaFunction,
    broadcast_col_binary_fn: CudaFunction,
    expand_fn: CudaFunction,
    transpose_2d_fn: CudaFunction,
    gather_fn: CudaFunction,
    gather_rows_fn: CudaFunction,
    gather_rows_u32_fn: CudaFunction,
    stream: Arc<CudaStream>,
}

impl ElementwiseKernels {
    /// Compile all element-wise kernels from a single CUDA source.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ELEMENTWISE_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("elementwise kernel compilation: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("elementwise module load: {e}"))
        })?;

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

        Ok(Self {
            binary_fn: load("elementwise_binary")?,
            unary_fn: load("elementwise_unary")?,
            broadcast_binary_fn: load("broadcast_scalar_binary")?,
            broadcast_row_binary_fn: load("broadcast_row_binary")?,
            broadcast_col_binary_fn: load("broadcast_col_binary")?,
            expand_fn: load("expand_broadcast")?,
            transpose_2d_fn: load("transpose_2d")?,
            gather_fn: load("gather_select")?,
            gather_rows_fn: load("gather_rows")?,
            gather_rows_u32_fn: load("gather_rows_u32")?,
            stream: Arc::clone(stream),
        })
    }

    /// Element-wise binary operation on two same-shape buffers.
    ///
    /// `op`: 0=add, 1=sub, 2=mul, 3=div, 4=min, 5=max
    pub fn binary(
        &self,
        a: &CudaSlice<f32>,
        b: &CudaSlice<f32>,
        n: usize,
        op: i32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
            MLError::ModelError(format!("binary alloc: {e}"))
        })?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* a, const float* b, float* out, int n, int op)
        // a and b are valid CudaSlice<f32> of at least n elements.
        // out is a freshly allocated CudaSlice<f32> of n elements.
        unsafe {
            self.stream
                .launch_builder(&self.binary_fn)
                .arg(a)
                .arg(b)
                .arg(&out)
                .arg(&n_i32)
                .arg(&op)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("elementwise_binary(op={op}): {e}")))?;
        }
        Ok(out)
    }

    /// Element-wise unary operation.
    ///
    /// `op`: 0=powf(param1), 1=sqr, 2=floor, 3=relu, 4=clamp(param1,param2),
    ///       5=abs, 6=neg, 7=exp, 8=log, 9=le(param1), 10=affine(v*param1+param2),
    ///       11=sin, 12=cos, 13=sqrt, 14=silu, 15=elu(param1=alpha), 16=recip,
    ///       17=sigmoid, 18=tanh
    pub fn unary(
        &self,
        x: &CudaSlice<f32>,
        n: usize,
        op: i32,
        param1: f32,
        param2: f32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
            MLError::ModelError(format!("unary alloc: {e}"))
        })?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* x, float* out, int n, int op, float param1, float param2)
        // x is a valid CudaSlice<f32> of at least n elements.
        // out is a freshly allocated CudaSlice<f32> of n elements.
        unsafe {
            self.stream
                .launch_builder(&self.unary_fn)
                .arg(x)
                .arg(&out)
                .arg(&n_i32)
                .arg(&op)
                .arg(&param1)
                .arg(&param2)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("elementwise_unary(op={op}): {e}")))?;
        }
        Ok(out)
    }

    /// Broadcast a scalar (1-element buffer) with a tensor via binary op.
    ///
    /// `op`: 0=mul, 1=div, 2=add, 3=sub
    /// If `scalar_is_lhs` is true: out[i] = scalar op tensor[i]
    /// Otherwise: out[i] = tensor[i] op scalar
    pub fn broadcast_scalar(
        &self,
        tensor: &CudaSlice<f32>,
        scalar: &CudaSlice<f32>,
        n: usize,
        op: i32,
        scalar_is_lhs: i32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
            MLError::ModelError(format!("broadcast_scalar alloc: {e}"))
        })?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* tensor, const float* scalar, float* out,
        //    int n, int op, int scalar_is_lhs)
        // tensor is at least n elements, scalar is at least 1 element.
        unsafe {
            self.stream
                .launch_builder(&self.broadcast_binary_fn)
                .arg(tensor)
                .arg(scalar)
                .arg(&out)
                .arg(&n_i32)
                .arg(&op)
                .arg(&scalar_is_lhs)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("broadcast_scalar(op={op}): {e}")))?;
        }
        Ok(out)
    }

    /// Row-broadcast binary op: [M, N] op [1, N] or [1, N] op [M, N].
    ///
    /// `op`: 0=mul, 1=div, 2=add, 3=sub
    /// `row_is_lhs`: if true, out[i] = row[i%cols] op matrix[i]; else matrix[i] op row[i%cols]
    pub fn broadcast_row(
        &self,
        matrix: &CudaSlice<f32>,
        row: &CudaSlice<f32>,
        total: usize,
        cols: usize,
        op: i32,
        row_is_lhs: i32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
            MLError::ModelError(format!("broadcast_row alloc: {e}"))
        })?;
        let total_i32 = total as i32;
        let cols_i32 = cols as i32;
        let cfg = elem_cfg(total);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* matrix, const float* row, float* out,
        //    int total, int cols, int op, int row_is_lhs)
        // matrix is at least `total` elements, row is at least `cols` elements.
        unsafe {
            self.stream
                .launch_builder(&self.broadcast_row_binary_fn)
                .arg(matrix)
                .arg(row)
                .arg(&out)
                .arg(&total_i32)
                .arg(&cols_i32)
                .arg(&op)
                .arg(&row_is_lhs)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("broadcast_row(op={op}): {e}")))?;
        }
        Ok(out)
    }

    /// Column-broadcast binary op: `[M, N]` op `[M, 1]` or `[M, 1]` op `[M, N]`.
    ///
    /// `op`: 0=mul, 1=div, 2=add, 3=sub
    /// `col_is_lhs`: if true, `out[i] = col_vec[i/cols] op matrix[i]`; else `matrix[i] op col_vec[i/cols]`
    pub fn broadcast_col(
        &self,
        matrix: &CudaSlice<f32>,
        col_vec: &CudaSlice<f32>,
        total: usize,
        cols: usize,
        op: i32,
        col_is_lhs: i32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
            MLError::ModelError(format!("broadcast_col alloc: {e}"))
        })?;
        let total_i32 = total as i32;
        let cols_i32 = cols as i32;
        let cfg = elem_cfg(total);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* matrix, const float* col_vec, float* out,
        //    int total, int cols, int op, int col_is_lhs)
        // matrix is at least `total` elements, col_vec is at least `total/cols` elements.
        unsafe {
            self.stream
                .launch_builder(&self.broadcast_col_binary_fn)
                .arg(matrix)
                .arg(col_vec)
                .arg(&out)
                .arg(&total_i32)
                .arg(&cols_i32)
                .arg(&op)
                .arg(&col_is_lhs)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("broadcast_col(op={op}): {e}")))?;
        }
        Ok(out)
    }

    /// Expand (broadcast) from source to target using pre-computed stride tables.
    ///
    /// `src_strides` and `src_shape` are GPU-resident u32 buffers of length `ndim`.
    /// `tgt_strides` is a GPU-resident u32 buffer of length `ndim`.
    pub fn expand(
        &self,
        src: &CudaSlice<f32>,
        src_strides_dev: &CudaSlice<u32>,
        src_shape_dev: &CudaSlice<u32>,
        tgt_strides_dev: &CudaSlice<u32>,
        total: usize,
        ndim: usize,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
            MLError::ModelError(format!("expand alloc: {e}"))
        })?;
        let total_i32 = total as i32;
        let ndim_i32 = ndim as i32;
        let cfg = elem_cfg(total);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* src, const unsigned int* src_strides,
        //    const unsigned int* src_shape, const unsigned int* tgt_strides,
        //    float* out, int total, int ndim)
        // All buffers are valid with sizes consistent with ndim and total.
        unsafe {
            self.stream
                .launch_builder(&self.expand_fn)
                .arg(src)
                .arg(src_strides_dev)
                .arg(src_shape_dev)
                .arg(tgt_strides_dev)
                .arg(&out)
                .arg(&total_i32)
                .arg(&ndim_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("expand_broadcast: {e}")))?;
        }
        Ok(out)
    }

    /// Transpose a 2-D matrix [rows, cols] -> [cols, rows].
    pub fn transpose_2d(
        &self,
        src: &CudaSlice<f32>,
        rows: usize,
        cols: usize,
    ) -> Result<CudaSlice<f32>, MLError> {
        let total = rows * cols;
        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
            MLError::ModelError(format!("transpose alloc: {e}"))
        })?;
        let total_i32 = total as i32;
        let rows_i32 = rows as i32;
        let cols_i32 = cols as i32;
        let cfg = elem_cfg(total);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* src, float* dst, int total, int rows, int cols)
        // src and dst are both valid CudaSlice<f32> of `total` elements.
        unsafe {
            self.stream
                .launch_builder(&self.transpose_2d_fn)
                .arg(src)
                .arg(&out)
                .arg(&total_i32)
                .arg(&rows_i32)
                .arg(&cols_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("transpose_2d: {e}")))?;
        }
        Ok(out)
    }

    /// Gather (index_select) along a dimension.
    ///
    /// For dim-0: out[i, ...] = src[indices[i], ...]
    /// `indices_dev` is a GPU-resident u32 buffer.
    pub fn gather(
        &self,
        src: &CudaSlice<f32>,
        indices_dev: &CudaSlice<u32>,
        num_indices: usize,
        inner_size: usize,
        outer_count: usize,
        src_dim_size: usize,
    ) -> Result<CudaSlice<f32>, MLError> {
        let total = outer_count * num_indices * inner_size;
        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
            MLError::ModelError(format!("gather alloc: {e}"))
        })?;
        let total_i32 = total as i32;
        let num_indices_i32 = num_indices as i32;
        let inner_size_i32 = inner_size as i32;
        let src_dim_stride_i32 = (src_dim_size * inner_size) as i32;
        let cfg = elem_cfg(total);

        // SAFETY: Kernel arguments match the CUDA function signature exactly:
        //   (const float* src, const unsigned int* indices, float* out,
        //    int total, int num_indices, int inner_size, int src_dim_stride)
        // src is valid with enough elements, indices has num_indices elements,
        // out has total elements allocated.
        unsafe {
            self.stream
                .launch_builder(&self.gather_fn)
                .arg(src)
                .arg(indices_dev)
                .arg(&out)
                .arg(&total_i32)
                .arg(&num_indices_i32)
                .arg(&inner_size_i32)
                .arg(&src_dim_stride_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("gather_select: {e}")))?;
        }
        Ok(out)
    }

    /// Per-row gather from a 2D matrix `[rows, cols]`.
    ///
    /// For each row `i`, picks the element at column `indices[i]` and writes
    /// it to `out[i]`. `indices` is a GPU-resident `i32` buffer of length `rows`.
    /// Output is a flat `CudaSlice<f32>` of length `rows`.
    ///
    /// This eliminates the need to download actions to CPU for per-row
    /// index selection from log-softmax matrices in PPO/policy gradient methods.
    pub fn gather_rows(
        &self,
        src: &CudaSlice<f32>,
        indices: &CudaSlice<i32>,
        rows: usize,
        cols: usize,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(rows).map_err(|e| {
            MLError::ModelError(format!("gather_rows alloc: {e}"))
        })?;
        let rows_i32 = rows as i32;
        let cols_i32 = cols as i32;
        let cfg = elem_cfg(rows);

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

    /// Per-row gather from a 2D matrix with `u32` indices.
    ///
    /// Same as [`gather_rows`](Self::gather_rows) but accepts `u32` index buffer
    /// (common for action indices stored as unsigned).
    pub fn gather_rows_u32(
        &self,
        src: &CudaSlice<f32>,
        indices: &CudaSlice<u32>,
        rows: usize,
        cols: usize,
    ) -> Result<CudaSlice<f32>, MLError> {
        let out = self.stream.alloc_zeros::<f32>(rows).map_err(|e| {
            MLError::ModelError(format!("gather_rows_u32 alloc: {e}"))
        })?;
        let rows_i32 = rows as i32;
        let cols_i32 = cols as i32;
        let cfg = elem_cfg(rows);

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

    // ── Convenience wrappers for unary ops 5-10 ───────────────────────

    /// Element-wise absolute value (`|x|`). Op 5.
    pub fn abs(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 5, 0.0, 0.0)
    }

    /// Element-wise negation (`-x`). Op 6.
    pub fn neg(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 6, 0.0, 0.0)
    }

    /// Element-wise exponential (`e^x`). Op 7.
    pub fn exp(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 7, 0.0, 0.0)
    }

    /// Element-wise natural logarithm (`ln(x)`). Op 8.
    pub fn log(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 8, 0.0, 0.0)
    }

    /// Element-wise less-or-equal comparison. Returns 1.0 if `x <= threshold`, else 0.0. Op 9.
    pub fn le(&self, x: &CudaSlice<f32>, threshold: f32, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 9, threshold, 0.0)
    }

    /// Element-wise affine transform: `x * mul + add`. Op 10.
    pub fn affine(&self, x: &CudaSlice<f32>, mul: f32, add: f32, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 10, mul, add)
    }

    /// Element-wise sine (`sin(x)`). Op 11.
    pub fn sin(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 11, 0.0, 0.0)
    }

    /// Element-wise cosine (`cos(x)`). Op 12.
    pub fn cos(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 12, 0.0, 0.0)
    }

    /// Element-wise square root (`sqrt(x)`). Op 13.
    pub fn sqrt(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 13, 0.0, 0.0)
    }

    /// Element-wise SiLU (`x * sigmoid(x)`). Op 14.
    pub fn silu(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 14, 0.0, 0.0)
    }

    /// Element-wise ELU (`x if x>=0, alpha*(exp(x)-1) otherwise`). Op 15.
    pub fn elu(&self, x: &CudaSlice<f32>, alpha: f32, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 15, alpha, 0.0)
    }

    /// Element-wise reciprocal (`1/x`, safe: returns 1e10 for `|x| < 1e-10`). Op 16.
    pub fn recip(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 16, 0.0, 0.0)
    }

    /// Element-wise sigmoid (`1 / (1 + exp(-x))`). Op 17.
    pub fn sigmoid(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 17, 0.0, 0.0)
    }

    /// Element-wise tanh. Op 18.
    pub fn tanh(&self, x: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.unary(x, n, 18, 0.0, 0.0)
    }

    /// Element-wise minimum of two same-shape buffers.
    pub fn min(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.binary(a, b, n, 4)
    }

    /// Element-wise maximum of two same-shape buffers.
    pub fn max(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, n: usize) -> Result<CudaSlice<f32>, MLError> {
        self.binary(a, b, n, 5)
    }

    /// Reference to the underlying stream.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }
}

/// Global singleton for compiled elementwise kernels.
///
/// Lazily compiled on first use. Thread-safe via `OnceLock`.
use std::collections::HashMap;
use parking_lot::Mutex;

/// Per-context kernel cache. Keyed by the Arc<CudaContext> pointer address,
/// so each distinct CudaContext gets its own compiled kernels.
/// This prevents stale CudaFunction handles when a context is released and
/// a new one is created (e.g., sequential GPU tests in the same process).
static KERNELS: Mutex<Option<HashMap<usize, Arc<ElementwiseKernels>>>> = Mutex::new(None);

/// Get or compile the elementwise kernels for the given stream's context.
///
/// Kernels are cached per CudaContext Arc identity (pointer address).
/// If a new CudaContext is created (different Arc), fresh kernels are compiled.
pub fn get_or_compile(stream: &Arc<CudaStream>) -> Result<Arc<ElementwiseKernels>, MLError> {
    // Use the Arc pointer address as key — unique per CudaContext instance
    let ctx_id = Arc::as_ptr(stream.context()) as usize;
    let mut cache = KERNELS.lock();
    let map = cache.get_or_insert_with(HashMap::new);

    if let Some(kernels) = map.get(&ctx_id) {
        return Ok(Arc::clone(kernels));
    }

    // Evict stale entries from previous contexts (they hold dead CudaFunction handles)
    map.clear();

    let kernels = Arc::new(ElementwiseKernels::new(stream)
        .map_err(|e| MLError::ModelError(format!("elementwise kernels: {e}")))?);
    map.insert(ctx_id, Arc::clone(&kernels));
    Ok(kernels)
}

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

const ELEMENTWISE_CUDA_SRC: &str = r#"
// ═══════════════════════════════════════════════════════════════════════
// Element-wise and memory-layout kernels for GpuTensor.
//
// All kernels use 1-D grid-stride pattern with 256 threads per block.
// ═══════════════════════════════════════════════════════════════════════

// ── Binary element-wise: add, sub, mul, div, min, max ────────────────

extern "C" __global__
void elementwise_binary(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ out,
    int n,
    int op
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float va = a[i];
        float vb = b[i];
        float r;
        switch (op) {
            case 0: r = va + vb; break;  // add
            case 1: r = va - vb; break;  // sub
            case 2: r = va * vb; break;  // mul
            case 3: r = va / vb; break;  // div
            case 4: r = fminf(va, vb); break;  // min
            case 5: r = fmaxf(va, vb); break;  // max
            default: r = 0.0f; break;
        }
        out[i] = r;
    }
}

// ── Unary element-wise: powf..affine, sin, cos, sqrt, silu, elu, recip ──

extern "C" __global__
void elementwise_unary(
    const float* __restrict__ x,
    float* __restrict__ out,
    int n,
    int op,
    float param1,
    float param2
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float v = x[i];
        float r;
        switch (op) {
            case 0: r = powf(v, param1); break;              // powf
            case 1: r = v * v; break;                         // sqr
            case 2: r = floorf(v); break;                     // floor
            case 3: r = fmaxf(v, 0.0f); break;               // relu
            case 4: r = fminf(fmaxf(v, param1), param2); break; // clamp
            case 5: r = fabsf(v); break;                         // abs
            case 6: r = -v; break;                                // neg
            case 7: r = expf(v); break;                           // exp
            case 8: r = logf(v); break;                           // log
            case 9: r = (v <= param1) ? 1.0f : 0.0f; break;      // le (compare)
            case 10: r = v * param1 + param2; break;              // affine (mul + add)
            case 11: r = sinf(v); break;                              // sin
            case 12: r = cosf(v); break;                              // cos
            case 13: r = sqrtf(v); break;                             // sqrt
            case 14: { float s = 1.0f / (1.0f + expf(-v)); r = v * s; break; } // silu
            case 15: r = (v >= 0.0f) ? v : param1 * (expf(v) - 1.0f); break;  // elu
            case 16: r = (fabsf(v) < 1e-10f) ? 1e10f : 1.0f / v; break;       // recip
            case 17: r = 1.0f / (1.0f + expf(-v)); break;                     // sigmoid
            case 18: r = tanhf(v); break;                                      // tanh
            default: r = v; break;
        }
        out[i] = r;
    }
}

// ── Broadcast scalar binary: tensor op scalar or scalar op tensor ────

extern "C" __global__
void broadcast_scalar_binary(
    const float* __restrict__ tensor,
    const float* __restrict__ scalar_ptr,
    float* __restrict__ out,
    int n,
    int op,
    int scalar_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float s = scalar_ptr[0];
        float t = tensor[i];
        float lhs = scalar_is_lhs ? s : t;
        float rhs = scalar_is_lhs ? t : s;
        float r;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = 0.0f; break;
        }
        out[i] = r;
    }
}

// ── Broadcast row binary: [M,N] op [1,N] or [1,N] op [M,N] ─────────

extern "C" __global__
void broadcast_row_binary(
    const float* __restrict__ matrix,
    const float* __restrict__ row,
    float* __restrict__ out,
    int total,
    int cols,
    int op,
    int row_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int col = i % cols;
        float m = matrix[i];
        float rv = row[col];
        float lhs = row_is_lhs ? rv : m;
        float rhs = row_is_lhs ? m  : rv;
        float r;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = 0.0f; break;
        }
        out[i] = r;
    }
}

// ── Broadcast column binary: [M,N] op [M,1] or [M,1] op [M,N] ──────

extern "C" __global__
void broadcast_col_binary(
    const float* __restrict__ matrix,
    const float* __restrict__ col_vec,
    float* __restrict__ out,
    int total,
    int cols,
    int op,
    int col_is_lhs
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int row = i / cols;
        float m = matrix[i];
        float cv = col_vec[row];
        float lhs = col_is_lhs ? cv : m;
        float rhs = col_is_lhs ? m  : cv;
        float r;
        switch (op) {
            case 0: r = lhs * rhs; break;  // mul
            case 1: r = lhs / rhs; break;  // div
            case 2: r = lhs + rhs; break;  // add
            case 3: r = lhs - rhs; break;  // sub
            default: r = 0.0f; break;
        }
        out[i] = r;
    }
}

// ── Expand (general N-D broadcast) ──────────────────────────────────

extern "C" __global__
void expand_broadcast(
    const float* __restrict__ src,
    const unsigned int* __restrict__ src_strides,
    const unsigned int* __restrict__ src_shape,
    const unsigned int* __restrict__ tgt_strides,
    float* __restrict__ out,
    int total,
    int ndim
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int src_idx = 0;
        int remaining = i;
        for (int d = 0; d < ndim; d++) {
            int tgt_stride = tgt_strides[d];
            int coord = remaining / tgt_stride;
            remaining = remaining % tgt_stride;
            int src_dim = src_shape[d];
            int src_coord = (src_dim == 1) ? 0 : coord;
            src_idx += src_coord * src_strides[d];
        }
        out[i] = src[src_idx];
    }
}

// ── Transpose 2-D ───────────────────────────────────────────────────

extern "C" __global__
void transpose_2d(
    const float* __restrict__ src,
    float* __restrict__ dst,
    int total,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int r = i / cols;
        int c = i % cols;
        // src[r, c] -> dst[c, r]
        dst[c * rows + r] = src[i];
    }
}

// ── Gather (index_select along a dimension) ──────────────────────────

extern "C" __global__
void gather_select(
    const float* __restrict__ src,
    const unsigned int* __restrict__ indices,
    float* __restrict__ out,
    int total,
    int num_indices,
    int inner_size,
    int src_dim_stride
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        // Decompose flat output index into (outer, idx_pos, inner)
        int out_idx_inner = i % inner_size;
        int tmp = i / inner_size;
        int idx_pos = tmp % num_indices;
        int outer = tmp / num_indices;

        unsigned int src_idx = indices[idx_pos];
        int src_flat = outer * src_dim_stride + (int)src_idx * inner_size + out_idx_inner;
        out[i] = src[src_flat];
    }
}

// ── Per-row gather: out[i] = src[i * cols + clamp(indices[i], 0, cols-1)] ──

extern "C" __global__
void gather_rows(
    const float* __restrict__ src,
    const int* __restrict__ indices,
    float* __restrict__ out,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < rows) {
        int col = indices[i];
        // Clamp to valid range [0, cols-1]
        col = col < 0 ? 0 : (col >= cols ? cols - 1 : col);
        out[i] = src[i * cols + col];
    }
}

// ── Per-row gather with unsigned int indices ────────────────────────

extern "C" __global__
void gather_rows_u32(
    const float* __restrict__ src,
    const unsigned int* __restrict__ indices,
    float* __restrict__ out,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < rows) {
        int col = (int)indices[i];
        col = col >= cols ? cols - 1 : col;
        out[i] = src[i * cols + col];
    }
}
"#;

#[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_binary_add() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let a_host = vec![1.0_f32, 2.0, 3.0, 4.0];
        let b_host = vec![10.0_f32, 20.0, 30.0, 40.0];
        let mut a = stream.alloc_zeros::<f32>(4).unwrap();
        let mut b = stream.alloc_zeros::<f32>(4).unwrap();
        stream.memcpy_htod(&a_host, &mut a).unwrap();
        stream.memcpy_htod(&b_host, &mut b).unwrap();

        let out = kernels.binary(&a, &b, 4, 0).unwrap(); // add
        let mut host = [0.0_f32; 4];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - 11.0).abs() < 1e-5);
        assert!((host[1] - 22.0).abs() < 1e-5);
        assert!((host[2] - 33.0).abs() < 1e-5);
        assert!((host[3] - 44.0).abs() < 1e-5);
    }

    #[test]
    fn test_binary_sub() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let a_host = vec![10.0_f32, 20.0, 30.0];
        let b_host = vec![1.0_f32, 2.0, 3.0];
        let mut a = stream.alloc_zeros::<f32>(3).unwrap();
        let mut b = stream.alloc_zeros::<f32>(3).unwrap();
        stream.memcpy_htod(&a_host, &mut a).unwrap();
        stream.memcpy_htod(&b_host, &mut b).unwrap();

        let out = kernels.binary(&a, &b, 3, 1).unwrap(); // sub
        let mut host = [0.0_f32; 3];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - 9.0).abs() < 1e-5);
        assert!((host[1] - 18.0).abs() < 1e-5);
        assert!((host[2] - 27.0).abs() < 1e-5);
    }

    #[test]
    fn test_binary_mul() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let a_host = vec![2.0_f32, 3.0, 4.0];
        let b_host = vec![5.0_f32, 6.0, 7.0];
        let mut a = stream.alloc_zeros::<f32>(3).unwrap();
        let mut b = stream.alloc_zeros::<f32>(3).unwrap();
        stream.memcpy_htod(&a_host, &mut a).unwrap();
        stream.memcpy_htod(&b_host, &mut b).unwrap();

        let out = kernels.binary(&a, &b, 3, 2).unwrap(); // mul
        let mut host = [0.0_f32; 3];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - 10.0).abs() < 1e-5);
        assert!((host[1] - 18.0).abs() < 1e-5);
        assert!((host[2] - 28.0).abs() < 1e-5);
    }

    #[test]
    fn test_unary_relu() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let x_host = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0];
        let mut x = stream.alloc_zeros::<f32>(5).unwrap();
        stream.memcpy_htod(&x_host, &mut x).unwrap();

        let out = kernels.unary(&x, 5, 3, 0.0, 0.0).unwrap(); // relu
        let mut host = [0.0_f32; 5];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0]).abs() < 1e-5);
        assert!((host[1]).abs() < 1e-5);
        assert!((host[2]).abs() < 1e-5);
        assert!((host[3] - 1.0).abs() < 1e-5);
        assert!((host[4] - 2.0).abs() < 1e-5);
    }

    #[test]
    fn test_unary_sqr() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let x_host = vec![1.0_f32, 2.0, 3.0, -4.0];
        let mut x = stream.alloc_zeros::<f32>(4).unwrap();
        stream.memcpy_htod(&x_host, &mut x).unwrap();

        let out = kernels.unary(&x, 4, 1, 0.0, 0.0).unwrap(); // sqr
        let mut host = [0.0_f32; 4];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - 1.0).abs() < 1e-5);
        assert!((host[1] - 4.0).abs() < 1e-5);
        assert!((host[2] - 9.0).abs() < 1e-5);
        assert!((host[3] - 16.0).abs() < 1e-5);
    }

    #[test]
    fn test_unary_clamp() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let x_host = vec![-5.0_f32, -1.0, 0.5, 2.0, 10.0];
        let mut x = stream.alloc_zeros::<f32>(5).unwrap();
        stream.memcpy_htod(&x_host, &mut x).unwrap();

        let out = kernels.unary(&x, 5, 4, -2.0, 3.0).unwrap(); // clamp(-2, 3)
        let mut host = [0.0_f32; 5];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - (-2.0)).abs() < 1e-5);
        assert!((host[1] - (-1.0)).abs() < 1e-5);
        assert!((host[2] - 0.5).abs() < 1e-5);
        assert!((host[3] - 2.0).abs() < 1e-5);
        assert!((host[4] - 3.0).abs() < 1e-5);
    }

    #[test]
    fn test_transpose_2d() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        // 2x3 matrix: [[1,2,3],[4,5,6]]
        let x_host = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        let mut x = stream.alloc_zeros::<f32>(6).unwrap();
        stream.memcpy_htod(&x_host, &mut x).unwrap();

        let out = kernels.transpose_2d(&x, 2, 3).unwrap();
        let mut host = [0.0_f32; 6];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        // Expected: [[1,4],[2,5],[3,6]] = [1,4,2,5,3,6]
        assert!((host[0] - 1.0).abs() < 1e-5);
        assert!((host[1] - 4.0).abs() < 1e-5);
        assert!((host[2] - 2.0).abs() < 1e-5);
        assert!((host[3] - 5.0).abs() < 1e-5);
        assert!((host[4] - 3.0).abs() < 1e-5);
        assert!((host[5] - 6.0).abs() < 1e-5);
    }

    #[test]
    fn test_broadcast_scalar_mul() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        let tensor_host = vec![1.0_f32, 2.0, 3.0, 4.0];
        let scalar_host = vec![10.0_f32];
        let mut tensor = stream.alloc_zeros::<f32>(4).unwrap();
        let mut scalar = stream.alloc_zeros::<f32>(1).unwrap();
        stream.memcpy_htod(&tensor_host, &mut tensor).unwrap();
        stream.memcpy_htod(&scalar_host, &mut scalar).unwrap();

        let out = kernels.broadcast_scalar(&tensor, &scalar, 4, 0, 0).unwrap(); // mul, scalar is rhs
        let mut host = [0.0_f32; 4];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        assert!((host[0] - 10.0).abs() < 1e-5);
        assert!((host[1] - 20.0).abs() < 1e-5);
        assert!((host[2] - 30.0).abs() < 1e-5);
        assert!((host[3] - 40.0).abs() < 1e-5);
    }

    #[test]
    fn test_gather_dim0() {
        let stream = make_stream();
        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
        // 4x3 src: row0=[0,1,2], row1=[3,4,5], row2=[6,7,8], row3=[9,10,11]
        let src_host: Vec<f32> = (0..12).map(|i| i as f32).collect();
        let mut src = stream.alloc_zeros::<f32>(12).unwrap();
        stream.memcpy_htod(&src_host, &mut src).unwrap();

        // Select rows 2, 0
        let idx_host = vec![2_u32, 0];
        let mut idx = stream.alloc_zeros::<u32>(2).unwrap();
        stream.memcpy_htod(&idx_host, &mut idx).unwrap();

        // dim=0, inner_size=3, outer_count=1, src_dim_size=4
        let out = kernels.gather(&src, &idx, 2, 3, 1, 4).unwrap();
        let mut host = [0.0_f32; 6];
        stream.memcpy_dtoh(&out, &mut host).unwrap(); // test-only readback
        // row2: [6,7,8], row0: [0,1,2]
        assert!((host[0] - 6.0).abs() < 1e-5);
        assert!((host[1] - 7.0).abs() < 1e-5);
        assert!((host[2] - 8.0).abs() < 1e-5);
        assert!((host[3] - 0.0).abs() < 1e-5);
        assert!((host[4] - 1.0).abs() < 1e-5);
        assert!((host[5] - 2.0).abs() < 1e-5);
    }
}
