//! Stream-carrying GPU tensor and linear layer + element-wise ops.
//!
//! Unlike [`super::GpuTensor`] (which takes `&Arc<CudaStream>` as a parameter),
//! the types here carry an `Arc<CudaStream>` internally. This makes them
//! convenient for model code where the stream is fixed per-layer.
//!
//! Used by `ml-supervised` models (TFT, Mamba-2, Liquid, KAN, xLSTM, Diffusion).
//!
//! All element-wise operations delegate to GPU-native CUDA kernels from
//! [`super::elementwise::ElementwiseKernels`]. No host roundtrips in hot paths.

#![allow(unsafe_code)] // Required for cuBLAS FFI and device pointer access

use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};

use crate::MLError;

// ── Pointer helpers ───────────────────────────────────────────────────────

/// Extract raw CUDA device pointer from a `CudaSlice` (read-only).
fn raw_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Extract raw mutable CUDA device pointer from a `CudaSlice`.
fn raw_ptr_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Get the cached elementwise kernels for a stream.
fn ew(stream: &Arc<CudaStream>) -> Result<Arc<super::elementwise::ElementwiseKernels>, MLError> {
    super::elementwise::get_or_compile(stream)
}

/// Async device-to-device copy with element offsets.
///
/// Copies `count` f32 elements from `src[src_offset..]` to `dst[dst_offset..]`.
fn dtod_copy(
    stream: &CudaStream,
    src: &CudaSlice<f32>,
    dst: &mut CudaSlice<f32>,
    src_offset: usize,
    dst_offset: usize,
    count: usize,
) -> Result<(), MLError> {
    if count == 0 {
        return Ok(());
    }
    let src_view = src.slice(src_offset..src_offset + count);
    let src_ptr = {
        let (ptr, guard) = src_view.device_ptr(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    // For dst we need a mutable pointer at dst_offset.
    // Get base ptr, then add offset in bytes.
    let dst_base = raw_ptr_mut(dst, stream);
    let dst_ptr = dst_base + (dst_offset * std::mem::size_of::<f32>()) as u64;
    let num_bytes = count * std::mem::size_of::<f32>();
    // SAFETY: Both src and dst point to valid device memory within their
    // respective CudaSlice allocations. The caller ensures offsets and counts
    // do not exceed buffer bounds. The copy is on the same CUDA context.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
        ).map_err(|e| MLError::DeviceError(format!("dtod_copy: {e}")))?;
    }
    Ok(())
}

// ── StreamTensor ──────────────────────────────────────────────────────────

/// A contiguous, row-major, f32 GPU tensor that carries its own stream handle.
///
/// Shape is stored alongside the `CudaSlice` so callers can reshape/validate
/// without a roundtrip to host.
#[derive(Clone)]
pub struct StreamTensor {
    /// Contiguous f32 data on device. Length = product of `shape`.
    pub data: CudaSlice<f32>,
    /// Row-major shape, e.g. `[batch, features]`.
    pub shape: Vec<usize>,
    /// CUDA stream handle.
    pub stream: Arc<CudaStream>,
}

impl std::fmt::Debug for StreamTensor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StreamTensor")
            .field("shape", &self.shape)
            .field("numel", &self.numel())
            .finish_non_exhaustive()
    }
}

impl StreamTensor {
    /// Total number of elements.
    pub fn numel(&self) -> usize {
        self.shape.iter().product()
    }

    /// Create a tensor filled with zeros.
    pub fn zeros(shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n: usize = shape.iter().product();
        let data = stream
            .alloc_zeros::<f32>(n)
            .map_err(|e| MLError::DeviceError(format!("StreamTensor zeros alloc: {e}")))?;
        Ok(Self {
            data,
            shape: shape.to_vec(),
            stream: Arc::clone(stream),
        })
    }

    /// Create a tensor from a host `Vec<f32>`.
    pub fn from_vec(values: Vec<f32>, shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n: usize = shape.iter().product();
        if values.len() != n {
            return Err(MLError::DimensionMismatch {
                expected: n,
                actual: values.len(),
            });
        }
        let data = stream
            .clone_htod(&values)
            .map_err(|e| MLError::DeviceError(format!("StreamTensor htod: {e}")))?;
        Ok(Self {
            data,
            shape: shape.to_vec(),
            stream: Arc::clone(stream),
        })
    }

    /// Create a tensor from random normal values (mu=0, sigma=`std_dev`).
    pub fn randn(shape: &[usize], std_dev: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        use rand::Rng;
        let n: usize = shape.iter().product();
        let mut rng = rand::thread_rng();
        let values: Vec<f32> = (0..n)
            .map(|_| {
                // Box-Muller transform for approximate normal distribution
                let u1: f64 = rng.gen_range(1e-10_f64..1.0);
                let u2: f64 = rng.gen_range(0.0_f64..std::f64::consts::TAU);
                let normal = (-2.0 * u1.ln()).sqrt() * u2.cos();
                (normal as f32) * std_dev
            })
            .collect();
        Self::from_vec(values, shape, stream)
    }

    /// Copy tensor data back to host.
    ///
    /// **WARNING**: This is a synchronous GPU-to-CPU download. It is NOT
    /// suitable for use in hot training / inference paths. Use it only for
    /// checkpointing, logging, or tests.
    #[cold]
    pub fn to_vec(&self) -> Result<Vec<f32>, MLError> {
        self.stream
            .clone_dtoh(&self.data) // gpu-exit: API boundary readback
            .map_err(|e| MLError::DeviceError(format!("StreamTensor dtoh: {e}")))
    }

    /// Reshape (no data copy -- just changes shape metadata).
    /// Total element count must match.
    pub fn reshape(&self, new_shape: &[usize]) -> Result<Self, MLError> {
        let new_n: usize = new_shape.iter().product();
        if new_n != self.numel() {
            return Err(MLError::DimensionMismatch {
                expected: self.numel(),
                actual: new_n,
            });
        }
        Ok(Self {
            data: self.data.clone(),
            shape: new_shape.to_vec(),
            stream: Arc::clone(&self.stream),
        })
    }

    /// Get a dimension by index (safe).
    pub fn dim(&self, idx: usize) -> Result<usize, MLError> {
        self.shape
            .get(idx)
            .copied()
            .ok_or_else(|| MLError::DimensionMismatch {
                expected: idx + 1,
                actual: self.shape.len(),
            })
    }
}

// ── StreamLinear ──────────────────────────────────────────────────────────

/// GPU-native linear layer: `y = x @ W^T + b`.
///
/// Weights are stored as `CudaSlice<f32>` and matrix multiplication is
/// dispatched via cuBLAS `sgemm`. Carries its own stream handle.
#[derive(Debug, Clone)]
pub struct StreamLinear {
    /// Weight matrix: shape (`out_features`, `in_features`) -- row-major.
    pub weight: CudaSlice<f32>,
    /// Bias vector: shape (`out_features`,).
    pub bias: Option<CudaSlice<f32>>,
    pub in_features: usize,
    pub out_features: usize,
    stream: Arc<CudaStream>,
}

impl StreamLinear {
    /// Create a new linear layer with Xavier-initialized weights and zero bias.
    pub fn new(
        in_features: usize,
        out_features: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        use rand::Rng;
        let limit = (6.0_f64 / (in_features + out_features) as f64).sqrt() as f32;
        let n = out_features * in_features;

        let mut rng = rand::thread_rng();
        let weight_data: Vec<f32> = (0..n).map(|_| rng.gen_range(-limit..limit)).collect();

        let weight = stream
            .clone_htod(&weight_data)
            .map_err(|e| MLError::DeviceError(format!("StreamLinear weight htod: {e}")))?;

        let bias_data = vec![0.0_f32; out_features];
        let bias = stream
            .clone_htod(&bias_data)
            .map_err(|e| MLError::DeviceError(format!("StreamLinear bias htod: {e}")))?;

        Ok(Self {
            weight,
            bias: Some(bias),
            in_features,
            out_features,
            stream: Arc::clone(stream),
        })
    }

    /// Create from existing weight/bias slices (e.g. loaded from checkpoint).
    pub fn from_slices(
        weight: CudaSlice<f32>,
        bias: Option<CudaSlice<f32>>,
        in_features: usize,
        out_features: usize,
        stream: &Arc<CudaStream>,
    ) -> Self {
        Self {
            weight,
            bias,
            in_features,
            out_features,
            stream: Arc::clone(stream),
        }
    }

    /// Forward pass: `y = x @ W^T + b`.
    ///
    /// `x` shape: `(M, in_features)`.
    /// Returns shape: `(M, out_features)`.
    pub fn forward(&self, x: &StreamTensor) -> Result<StreamTensor, MLError> {
        let ndim = x.shape.len();
        if ndim < 2 {
            return Err(MLError::DimensionMismatch {
                expected: 2,
                actual: ndim,
            });
        }

        let k = x.dim(ndim - 1)?;
        if k != self.in_features {
            return Err(MLError::DimensionMismatch {
                expected: self.in_features,
                actual: k,
            });
        }

        let m: usize = x.shape.iter().take(ndim - 1).product();
        let n = self.out_features;

        let mut out = self
            .stream
            .alloc_zeros::<f32>(m * n)
            .map_err(|e| MLError::DeviceError(format!("StreamLinear output alloc: {e}")))?;

        let blas = cudarc::cublas::CudaBlas::new(Arc::clone(&self.stream))
            .map_err(|e| MLError::DeviceError(format!("cuBLAS init: {e}")))?;

        let w_ptr = raw_ptr(&self.weight, &self.stream);
        let x_ptr = raw_ptr(&x.data, &self.stream);
        let y_ptr = raw_ptr_mut(&mut out, &self.stream);

        // SAFETY: cuBLAS sgemm operates on valid device pointers from pre-allocated
        // weight [n,k], input [m,k], and output [m,n] CudaSlice buffers. All
        // dimensions are validated from StreamLinear construction parameters.
        unsafe {
            cudarc::cublas::result::sgemm(
                *blas.handle(),
                cublasOperation_t::CUBLAS_OP_T,
                cublasOperation_t::CUBLAS_OP_N,
                n as i32,
                m as i32,
                k as i32,
                &1.0_f32 as *const f32,
                w_ptr as *const f32,
                k as i32,
                x_ptr as *const f32,
                k as i32,
                &0.0_f32 as *const f32,
                y_ptr as *mut f32,
                n as i32,
            )
            .map_err(|e| MLError::DeviceError(format!("cuBLAS sgemm: {e}")))?;
        }

        // Add bias via GPU broadcast_row (row = bias[1,N], matrix = out[M,N])
        if let Some(ref bias) = self.bias {
            let kernels = ew(&self.stream)?;
            // broadcast_row: matrix[M,N] + row[1,N] => out[M,N], op=2 (add), row_is_lhs=0
            out = kernels.broadcast_row(&out, bias, m * n, n, 2, 0)?;
        }

        let mut out_shape = x.shape.clone();
        if let Some(last) = out_shape.last_mut() {
            *last = n;
        }

        Ok(StreamTensor {
            data: out,
            shape: out_shape,
            stream: Arc::clone(&self.stream),
        })
    }

    /// Collect weight data to host for checkpoint serialization.
    #[cold]
    pub fn weight_to_vec(&self) -> Result<Vec<f32>, MLError> {
        self.stream
            .clone_dtoh(&self.weight) // gpu-exit: checkpoint weight export
            .map_err(|e| MLError::DeviceError(format!("weight dtoh: {e}")))
    }

    /// Collect bias data to host for checkpoint serialization.
    #[cold]
    pub fn bias_to_vec(&self) -> Result<Option<Vec<f32>>, MLError> {
        match &self.bias {
            Some(b) => {
                let v = self
                    .stream
                    .clone_dtoh(b) // gpu-exit: checkpoint weight export
                    .map_err(|e| MLError::DeviceError(format!("bias dtoh: {e}")))?;
                Ok(Some(v))
            }
            None => Ok(None),
        }
    }
}

// ── Element-wise ops (GPU-native via ElementwiseKernels) ─────────────────

/// Helper: wrap a `CudaSlice` result from a kernel into a `StreamTensor`.
#[inline]
fn wrap(data: CudaSlice<f32>, shape: &[usize], stream: &Arc<CudaStream>) -> StreamTensor {
    StreamTensor {
        data,
        shape: shape.to_vec(),
        stream: Arc::clone(stream),
    }
}

/// Element-wise `SiLU` (x * sigmoid(x)) on a `StreamTensor`.
pub fn gpu_silu(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.silu(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise sigmoid on a `StreamTensor`.
pub fn gpu_sigmoid(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.sigmoid(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise tanh on a `StreamTensor`.
pub fn gpu_tanh(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.tanh(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise add: `a + b` (must have same shape).
pub fn gpu_add(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape != b.shape {
        return Err(MLError::DimensionMismatch {
            expected: a.numel(),
            actual: b.numel(),
        });
    }
    let k = ew(&a.stream)?;
    let out = k.binary(&a.data, &b.data, a.numel(), 0)?; // 0 = add
    Ok(wrap(out, &a.shape, &a.stream))
}

/// Element-wise subtract: `a - b` (must have same shape).
pub fn gpu_sub(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape != b.shape {
        return Err(MLError::DimensionMismatch {
            expected: a.numel(),
            actual: b.numel(),
        });
    }
    let k = ew(&a.stream)?;
    let out = k.binary(&a.data, &b.data, a.numel(), 1)?; // 1 = sub
    Ok(wrap(out, &a.shape, &a.stream))
}

/// Element-wise multiply: `a * b` (must have same shape).
pub fn gpu_mul(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape != b.shape {
        return Err(MLError::DimensionMismatch {
            expected: a.numel(),
            actual: b.numel(),
        });
    }
    let k = ew(&a.stream)?;
    let out = k.binary(&a.data, &b.data, a.numel(), 2)?; // 2 = mul
    Ok(wrap(out, &a.shape, &a.stream))
}

/// Scalar multiply: `x * scalar`.
pub fn gpu_scale(x: &StreamTensor, scalar: f32) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.affine(&x.data, scalar, 0.0, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise exp.
pub fn gpu_exp(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.exp(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise reciprocal: 1/x.
pub fn gpu_recip(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.recip(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Scalar add: x + scalar (broadcast).
pub fn gpu_add_scalar(x: &StreamTensor, scalar: f32) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.affine(&x.data, 1.0, scalar, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Mean over all elements, returns a single f32.
pub fn gpu_mean_all(x: &StreamTensor) -> Result<f32, MLError> {
    let n = x.numel();
    if n == 0 {
        return Err(MLError::InvalidInput("Cannot compute mean of empty tensor".to_owned()));
    }
    let red = super::reductions::ReductionKernels::new(&x.stream)?;
    let stats = red.stats(&x.data, n)?;
    Ok(stats.mean)
}

/// Concatenate two 2D tensors along dimension 1 (columns).
///
/// Uses GPU DtoD copies (narrow slices) -- no host roundtrip.
pub fn gpu_cat_dim1(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    let ndim = a.shape.len();
    if ndim != b.shape.len() {
        return Err(MLError::InvalidInput(format!(
            "gpu_cat_dim1: rank mismatch {} vs {}",
            ndim,
            b.shape.len()
        )));
    }
    if ndim < 2 || ndim > 3 {
        return Err(MLError::InvalidInput(format!(
            "gpu_cat_dim1 requires 2D or 3D tensors, got {ndim}D"
        )));
    }

    // Dim-0 must match (batch for 2D; batch for 3D)
    let dim0_a = a.dim(0)?;
    let dim0_b = b.dim(0)?;
    if dim0_a != dim0_b {
        return Err(MLError::DimensionMismatch {
            expected: dim0_a,
            actual: dim0_b,
        });
    }

    // For 3D, dim-2 must also match (features)
    if ndim == 3 {
        let dim2_a = a.dim(2)?;
        let dim2_b = b.dim(2)?;
        if dim2_a != dim2_b {
            return Err(MLError::DimensionMismatch {
                expected: dim2_a,
                actual: dim2_b,
            });
        }
    }

    let dim1_a = a.dim(1)?;
    let dim1_b = b.dim(1)?;
    let dim1_out = dim1_a + dim1_b;

    if ndim == 2 {
        // 2D: [rows, cols] concat along dim-1 (columns)
        // Per-row DtoD copy of a's cols then b's cols into output
        let rows = dim0_a;
        let cols_out = dim1_out;
        let total = rows * cols_out;
        let mut out = a.stream.alloc_zeros::<f32>(total)
            .map_err(|e| MLError::DeviceError(format!("cat_dim1 alloc: {e}")))?;

        for r in 0..rows {
            let a_offset = r * dim1_a;
            let b_offset = r * dim1_b;
            let o_offset = r * cols_out;
            // Copy a's row segment
            dtod_copy(&a.stream, &a.data, &mut out, a_offset, o_offset, dim1_a)?;
            // Copy b's row segment
            dtod_copy(&a.stream, &b.data, &mut out, b_offset, o_offset + dim1_a, dim1_b)?;
        }
        Ok(StreamTensor {
            data: out,
            shape: vec![rows, cols_out],
            stream: Arc::clone(&a.stream),
        })
    } else {
        // 3D: [batch, seq, features] concat along dim-1 (sequence)
        let batch = dim0_a;
        let seq_a = dim1_a;
        let seq_b = dim1_b;
        let seq_out = dim1_out;
        let feat = a.dim(2)?;

        let stride_a = seq_a * feat;
        let stride_b = seq_b * feat;
        let stride_out = seq_out * feat;
        let total = batch * stride_out;

        let mut out = a.stream.alloc_zeros::<f32>(total)
            .map_err(|e| MLError::DeviceError(format!("cat_dim1_3d alloc: {e}")))?;

        for bi in 0..batch {
            // Copy a's block
            dtod_copy(&a.stream, &a.data, &mut out, bi * stride_a, bi * stride_out, stride_a)?;
            // Copy b's block
            dtod_copy(&a.stream, &b.data, &mut out, bi * stride_b, bi * stride_out + stride_a, stride_b)?;
        }
        Ok(StreamTensor {
            data: out,
            shape: vec![batch, seq_out, feat],
            stream: Arc::clone(&a.stream),
        })
    }
}

/// Matmul for 2D tensors: C(M,N) = A(M,K) * B(K,N).
pub fn gpu_matmul(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape.len() != 2 || b.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_matmul requires 2D tensors".to_owned()));
    }
    let m = a.dim(0)?;
    let k_a = a.dim(1)?;
    let k_b = b.dim(0)?;
    let n = b.dim(1)?;

    if k_a != k_b {
        return Err(MLError::DimensionMismatch {
            expected: k_a,
            actual: k_b,
        });
    }
    let k = k_a;

    let mut out = a
        .stream
        .alloc_zeros::<f32>(m * n)
        .map_err(|e| MLError::DeviceError(format!("matmul output alloc: {e}")))?;

    let blas = cudarc::cublas::CudaBlas::new(Arc::clone(&a.stream))
        .map_err(|e| MLError::DeviceError(format!("cuBLAS init: {e}")))?;

    // C_col(N,M) = B_col(N,K) * A_col(K,M)
    let b_ptr = raw_ptr(&b.data, &a.stream);
    let a_ptr = raw_ptr(&a.data, &a.stream);
    let c_ptr = raw_ptr_mut(&mut out, &a.stream);

    // SAFETY: cuBLAS sgemm computes C = B * A in column-major. All device
    // pointers are valid from pre-allocated CudaSlice buffers with compatible
    // matrix dimensions (M, K, N derived from validated tensor shapes).
    unsafe {
        cudarc::cublas::result::sgemm(
            *blas.handle(),
            cublasOperation_t::CUBLAS_OP_N,
            cublasOperation_t::CUBLAS_OP_N,
            n as i32,
            m as i32,
            k as i32,
            &1.0_f32 as *const f32,
            b_ptr as *const f32,
            n as i32,
            a_ptr as *const f32,
            k as i32,
            &0.0_f32 as *const f32,
            c_ptr as *mut f32,
            n as i32,
        )
        .map_err(|e| MLError::DeviceError(format!("cuBLAS sgemm matmul: {e}")))?;
    }

    Ok(StreamTensor {
        data: out,
        shape: vec![m, n],
        stream: Arc::clone(&a.stream),
    })
}

/// Element-wise affine transform: `x * scale + offset`.
pub fn gpu_affine(x: &StreamTensor, scale: f64, offset: f64) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.affine(&x.data, scale as f32, offset as f32, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Flatten a tensor from dim `start` to `end` (inclusive).
pub fn gpu_flatten(x: &StreamTensor, start: usize, end: usize) -> Result<StreamTensor, MLError> {
    if start >= x.shape.len() || end >= x.shape.len() || start > end {
        return Err(MLError::InvalidInput(format!(
            "Invalid flatten range [{}, {}] for shape {:?}",
            start, end, x.shape
        )));
    }
    let mut new_shape = Vec::new();
    for (i, &d) in x.shape.iter().enumerate() {
        if i < start || i > end {
            new_shape.push(d);
        } else if i == start {
            let flat: usize = x.shape.get(start..=end).unwrap_or(&[]).iter().product();
            new_shape.push(flat);
        } else {
            // Dimensions between start+1..end are consumed by the flatten above
        }
    }
    x.reshape(&new_shape)
}

/// Extract a slice along dim 1 for a 3D tensor: `x[:, idx, :]`.
/// Returns a 2D tensor of shape `(batch, features)`.
///
/// Uses GPU DtoD narrow copy -- no host roundtrip.
pub fn gpu_select_dim1(x: &StreamTensor, idx: usize) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 3 {
        return Err(MLError::InvalidInput("gpu_select_dim1 requires 3D tensor".to_owned()));
    }
    let batch = x.dim(0)?;
    let seq_len = x.dim(1)?;
    let features = x.dim(2)?;
    if idx >= seq_len {
        return Err(MLError::DimensionMismatch {
            expected: seq_len,
            actual: idx + 1,
        });
    }

    let out_elems = batch * features;
    let mut out = x.stream.alloc_zeros::<f32>(out_elems)
        .map_err(|e| MLError::DeviceError(format!("select_dim1 alloc: {e}")))?;

    for bi in 0..batch {
        let src_offset = bi * seq_len * features + idx * features;
        let dst_offset = bi * features;
        dtod_copy(&x.stream, &x.data, &mut out, src_offset, dst_offset, features)?;
    }

    Ok(StreamTensor {
        data: out,
        shape: vec![batch, features],
        stream: Arc::clone(&x.stream),
    })
}

/// Element-wise square: x^2.
pub fn gpu_sqr(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.unary(&x.data, x.numel(), 1, 0.0, 0.0)?; // 1 = sqr
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise sin.
pub fn gpu_sin(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.sin(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise cos.
pub fn gpu_cos(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.cos(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Scalar subtraction: scalar - x (element-wise).
pub fn gpu_scalar_sub(scalar: f32, x: &StreamTensor) -> Result<StreamTensor, MLError> {
    // scalar - x = -x + scalar = x * (-1) + scalar
    let k = ew(&x.stream)?;
    let out = k.affine(&x.data, -1.0, scalar, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Create a filled tensor (all elements = value).
pub fn gpu_full(shape: &[usize], value: f32, stream: &Arc<CudaStream>) -> Result<StreamTensor, MLError> {
    let n: usize = shape.iter().product();
    // Create zeros then use affine(mul=0, add=value) to fill
    let zeros = stream.alloc_zeros::<f32>(n)
        .map_err(|e| MLError::DeviceError(format!("gpu_full alloc: {e}")))?;
    let k = ew(stream)?;
    let out = k.affine(&zeros, 0.0, value, n)?;
    Ok(wrap(out, shape, stream))
}

/// Gather along dimension 0 of a 2D source using 2D float indices.
///
/// This uses a host roundtrip for the index mapping since the indices are
/// float-typed and require clamping logic not suited for the gather kernel.
pub fn gpu_gather_dim0(source: &StreamTensor, indices: &StreamTensor) -> Result<StreamTensor, MLError> {
    if source.shape.len() != 2 || indices.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_gather_dim0 requires 2D tensors".to_owned()));
    }
    let s_rows = source.dim(0)?;
    let s_cols = source.dim(1)?;
    let n = indices.dim(0)?;
    let i_cols = indices.dim(1)?;
    if s_cols != i_cols {
        return Err(MLError::DimensionMismatch {
            expected: s_cols,
            actual: i_cols,
        });
    }

    // Download only the indices (small), keep source on GPU
    let idx_host = indices.to_vec()?; // gpu-exit: small index array for gather
    let out_elems = n * s_cols;
    let mut out = source.stream.alloc_zeros::<f32>(out_elems)
        .map_err(|e| MLError::DeviceError(format!("gather_dim0 alloc: {e}")))?;

    // For each row in output, copy the corresponding source row via DtoD
    for i in 0..n {
        for j in 0..s_cols {
            let idx_val = idx_host.get(i * s_cols + j).copied().unwrap_or(0.0) as usize;
            let clamped_idx = idx_val.min(s_rows.saturating_sub(1));
            // Single-element DtoD copy
            dtod_copy(&source.stream, &source.data, &mut out, clamped_idx * s_cols + j, i * s_cols + j, 1)?;
        }
    }

    Ok(StreamTensor {
        data: out,
        shape: vec![n, s_cols],
        stream: Arc::clone(&source.stream),
    })
}

/// Broadcast a 1D (N,) tensor to (N, C) by repeating along dim 1.
///
/// Uses GPU broadcast_col kernel: treats x as [N,1] column vector,
/// broadcasts to [N, new_cols] via identity multiplication with ones.
pub fn gpu_unsqueeze_broadcast(x: &StreamTensor, new_cols: usize) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 1 && !(x.shape.len() == 2 && x.dim(1)? == 1) {
        return Err(MLError::InvalidInput(format!(
            "gpu_unsqueeze_broadcast expects 1D or (N,1), got {:?}",
            x.shape
        )));
    }
    let n = x.dim(0)?;
    let total = n * new_cols;

    // Create a ones matrix [N, new_cols], then broadcast_col multiply with x[N,1]
    // col_is_lhs=1, op=0(mul): out[i] = x[i/cols] * ones[i] = x[i/cols]
    // Actually simpler: use affine on a zeros buffer, but we need broadcast.
    // Use broadcast_col with ones: matrix=ones[N,C], col=x[N,1], op=0(mul), col_is_lhs=1
    // => out[i] = x[i/cols] * 1.0 = x[i/cols]
    let ones = x.stream.alloc_zeros::<f32>(total)
        .map_err(|e| MLError::DeviceError(format!("unsqueeze_broadcast ones alloc: {e}")))?;
    // Fill ones using affine(mul=0, add=1)
    let k = ew(&x.stream)?;
    let ones_filled = k.affine(&ones, 0.0, 1.0, total)?;
    // broadcast_col: matrix[N,C] op col[N,1], op=0(mul), col_is_lhs=1
    // => out[i] = x[i/new_cols] * ones_filled[i]
    let out = k.broadcast_col(&ones_filled, &x.data, total, new_cols, 0, 1)?;
    Ok(wrap(out, &[n, new_cols], &x.stream))
}

/// Element-wise clamp.
pub fn gpu_clamp(x: &StreamTensor, min: f32, max: f32) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.unary(&x.data, x.numel(), 4, min, max)?; // 4 = clamp
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise floor.
pub fn gpu_floor(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.unary(&x.data, x.numel(), 2, 0.0, 0.0)?; // 2 = floor
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise minimum of two tensors.
pub fn gpu_minimum(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape != b.shape {
        return Err(MLError::DimensionMismatch {
            expected: a.numel(),
            actual: b.numel(),
        });
    }
    let k = ew(&a.stream)?;
    let out = k.min(&a.data, &b.data, a.numel())?;
    Ok(wrap(out, &a.shape, &a.stream))
}

/// Sum across dimension `dim` for a 2D tensor.
pub fn gpu_sum_dim(x: &StreamTensor, dim: usize) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_sum_dim only supports 2D tensors".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;

    if dim == 1 {
        // Sum across columns for each row -> [rows, 1]
        // For each row, use reduction kernel on that row's slice.
        // Since ReductionKernels::sum operates on contiguous data, we can
        // process each row as a contiguous slice.
        let red = super::reductions::ReductionKernels::new(&x.stream)?;
        let mut row_sums = Vec::with_capacity(rows);
        for r in 0..rows {
            // Narrow to row r: offset = r * cols, len = cols
            let mut row_buf = x.stream.alloc_zeros::<f32>(cols)
                .map_err(|e| MLError::DeviceError(format!("sum_dim row alloc: {e}")))?;
            dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?;
            let s = red.sum(&row_buf, cols)?;
            row_sums.push(s);
        }
        StreamTensor::from_vec(row_sums, &[rows, 1], &x.stream)
    } else if dim == 0 {
        // Sum across rows for each column -> [1, cols]
        let red = super::reductions::ReductionKernels::new(&x.stream)?;
        let col_sums_host = red.col_sums(&x.data, rows, cols)?;
        StreamTensor::from_vec(col_sums_host, &[1, cols], &x.stream)
    } else {
        Err(MLError::InvalidInput(format!("Invalid dim {dim} for 2D tensor")))
    }
}

/// Element-wise `ReLU`: max(0, x).
pub fn gpu_relu(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.unary(&x.data, x.numel(), 3, 0.0, 0.0)?; // 3 = relu
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise ELU: x if x >= 0, alpha*(exp(x)-1) otherwise.
pub fn gpu_elu(x: &StreamTensor, alpha: f32) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.elu(&x.data, alpha, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise negation: -x.
pub fn gpu_neg(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.neg(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise square root.
pub fn gpu_sqrt(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.sqrt(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise log (natural logarithm).
pub fn gpu_log(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.log(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise absolute value.
pub fn gpu_abs(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    let k = ew(&x.stream)?;
    let out = k.abs(&x.data, x.numel())?;
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Narrow (slice) along a dimension for a 2D tensor.
///
/// `gpu_narrow(x, dim, start, len)` extracts a sub-tensor along `dim`.
/// Uses GPU DtoD copies -- no host roundtrip.
pub fn gpu_narrow_2d(x: &StreamTensor, dim: usize, start: usize, len: usize) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_narrow_2d requires 2D tensor".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;

    if dim == 0 {
        if start + len > rows {
            return Err(MLError::DimensionMismatch { expected: rows, actual: start + len });
        }
        // Contiguous slice: just copy from offset
        let out_elems = len * cols;
        let mut out = x.stream.alloc_zeros::<f32>(out_elems)
            .map_err(|e| MLError::DeviceError(format!("narrow_2d alloc: {e}")))?;
        dtod_copy(&x.stream, &x.data, &mut out, start * cols, 0, out_elems)?;
        Ok(StreamTensor {
            data: out,
            shape: vec![len, cols],
            stream: Arc::clone(&x.stream),
        })
    } else if dim == 1 {
        if start + len > cols {
            return Err(MLError::DimensionMismatch { expected: cols, actual: start + len });
        }
        // Non-contiguous: copy per-row
        let out_elems = rows * len;
        let mut out = x.stream.alloc_zeros::<f32>(out_elems)
            .map_err(|e| MLError::DeviceError(format!("narrow_2d_dim1 alloc: {e}")))?;
        for r in 0..rows {
            dtod_copy(&x.stream, &x.data, &mut out, r * cols + start, r * len, len)?;
        }
        Ok(StreamTensor {
            data: out,
            shape: vec![rows, len],
            stream: Arc::clone(&x.stream),
        })
    } else {
        Err(MLError::InvalidInput(format!("Invalid dim {dim} for 2D narrow")))
    }
}

/// Softmax along the last dimension of a 2D tensor.
///
/// Uses GPU kernels for max, subtract, exp, and sum -- no host roundtrip.
pub fn gpu_softmax(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_softmax requires 2D tensor".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;
    let k = ew(&x.stream)?;
    let red = super::reductions::ReductionKernels::new(&x.stream)?;

    // Per-row softmax: for each row, subtract max, exp, divide by sum.
    // Process row by row to use contiguous reduction.
    let total = rows * cols;
    let mut out = x.stream.alloc_zeros::<f32>(total)
        .map_err(|e| MLError::DeviceError(format!("softmax alloc: {e}")))?;

    for r in 0..rows {
        // Extract row
        let mut row_buf = x.stream.alloc_zeros::<f32>(cols)
            .map_err(|e| MLError::DeviceError(format!("softmax row alloc: {e}")))?;
        dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?;

        // Find max for numerical stability
        let stats = red.stats(&row_buf, cols)?;
        let max_val = stats.max;

        // Subtract max: affine(mul=1, add=-max)
        let shifted = k.affine(&row_buf, 1.0, -max_val, cols)?;

        // Exp
        let exp_row = k.exp(&shifted, cols)?;

        // Sum
        let sum_val = red.sum(&exp_row, cols)?;

        // Divide by sum: affine(mul=1/sum, add=0)
        let inv_sum = if sum_val > 0.0 { 1.0 / sum_val } else { 0.0 };
        let normalized = k.affine(&exp_row, inv_sum, 0.0, cols)?;

        // Copy back to output
        dtod_copy(&x.stream, &normalized, &mut out, 0, r * cols, cols)?;
    }

    Ok(StreamTensor {
        data: out,
        shape: vec![rows, cols],
        stream: Arc::clone(&x.stream),
    })
}

/// Layer normalization on the last dimension of a 2D tensor.
///
/// Normalizes each row to zero mean and unit variance, then applies
/// affine transform: `y = gamma * (x - mean) / sqrt(var + eps) + beta`.
///
/// Uses GPU kernels for stats, subtract, multiply, add -- no host roundtrip.
pub fn gpu_layer_norm(
    x: &StreamTensor,
    gamma: &StreamTensor,
    beta: &StreamTensor,
    eps: f32,
) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_layer_norm requires 2D tensor".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;
    let k = ew(&x.stream)?;
    let red = super::reductions::ReductionKernels::new(&x.stream)?;

    let total = rows * cols;
    let mut out = x.stream.alloc_zeros::<f32>(total)
        .map_err(|e| MLError::DeviceError(format!("layer_norm alloc: {e}")))?;

    for r in 0..rows {
        // Extract row
        let mut row_buf = x.stream.alloc_zeros::<f32>(cols)
            .map_err(|e| MLError::DeviceError(format!("layer_norm row alloc: {e}")))?;
        dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?;

        // Compute mean and variance
        let stats = red.stats(&row_buf, cols)?;
        let mean = stats.mean;
        let var = stats.variance;
        let inv_std = 1.0 / (var + eps).sqrt();

        // Normalize: (x - mean) * inv_std = x * inv_std + (-mean * inv_std)
        let normalized = k.affine(&row_buf, inv_std, -mean * inv_std, cols)?;

        // Apply gamma and beta: result = normalized * gamma + beta
        // This is element-wise mul then add with gamma and beta vectors
        let scaled = k.binary(&normalized, &gamma.data, cols, 2)?; // mul
        let shifted = k.binary(&scaled, &beta.data, cols, 0)?; // add

        // Copy back
        dtod_copy(&x.stream, &shifted, &mut out, 0, r * cols, cols)?;
    }

    Ok(StreamTensor {
        data: out,
        shape: vec![rows, cols],
        stream: Arc::clone(&x.stream),
    })
}

/// Mean across last dimension of a 2D tensor, keeping dim as size 1.
pub fn gpu_mean_keepdim_last(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_mean_keepdim_last requires 2D tensor".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;
    let red = super::reductions::ReductionKernels::new(&x.stream)?;

    let mut result_host = Vec::with_capacity(rows);
    for r in 0..rows {
        let mut row_buf = x.stream.alloc_zeros::<f32>(cols)
            .map_err(|e| MLError::DeviceError(format!("mean_keepdim row alloc: {e}")))?;
        dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?;
        let stats = red.stats(&row_buf, cols)?;
        result_host.push(stats.mean);
    }
    StreamTensor::from_vec(result_host, &[rows, 1], &x.stream)
}

/// Broadcast multiply: a (rows, cols) * b (rows, 1) -> (rows, cols).
/// b is broadcast along columns.
pub fn gpu_broadcast_mul_col(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape.len() != 2 || b.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_broadcast_mul_col requires 2D tensors".to_owned()));
    }
    let rows = a.dim(0)?;
    let cols = a.dim(1)?;
    if b.dim(0)? != rows || b.dim(1)? != 1 {
        return Err(MLError::DimensionMismatch { expected: rows, actual: b.dim(0)? });
    }
    let k = ew(&a.stream)?;
    let total = rows * cols;
    // broadcast_col: matrix[M,N] op col[M,1], op=0(mul), col_is_lhs=0
    // => out[i] = matrix[i] * col[i/cols]
    let out = k.broadcast_col(&a.data, &b.data, total, cols, 0, 0)?;
    Ok(wrap(out, &[rows, cols], &a.stream))
}

/// Transpose a 2D tensor.
pub fn gpu_transpose(x: &StreamTensor) -> Result<StreamTensor, MLError> {
    if x.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_transpose requires 2D tensor".to_owned()));
    }
    let rows = x.dim(0)?;
    let cols = x.dim(1)?;
    let k = ew(&x.stream)?;
    let out = k.transpose_2d(&x.data, rows, cols)?;
    Ok(wrap(out, &[cols, rows], &x.stream))
}

/// Stack a list of 2D tensors into a 3D tensor along a new first dimension.
///
/// Input: N tensors each of shape (R, C).
/// Output: (N, R, C).
///
/// Uses GPU DtoD copies -- no host roundtrip.
pub fn gpu_stack_2d(tensors: &[StreamTensor]) -> Result<StreamTensor, MLError> {
    if tensors.is_empty() {
        return Err(MLError::InvalidInput("gpu_stack_2d: empty list".to_owned()));
    }
    let first = tensors.first().ok_or_else(|| MLError::InvalidInput("empty".to_owned()))?;
    let rows = first.dim(0)?;
    let cols = first.dim(1)?;
    let n = tensors.len();
    let elem_per_tensor = rows * cols;
    let total = n * elem_per_tensor;

    let mut out = first.stream.alloc_zeros::<f32>(total)
        .map_err(|e| MLError::DeviceError(format!("stack_2d alloc: {e}")))?;

    for (i, t) in tensors.iter().enumerate() {
        if t.dim(0)? != rows || t.dim(1)? != cols {
            return Err(MLError::DimensionMismatch { expected: rows * cols, actual: t.numel() });
        }
        dtod_copy(&first.stream, &t.data, &mut out, 0, i * elem_per_tensor, elem_per_tensor)?;
    }

    let stream = &first.stream;
    Ok(StreamTensor {
        data: out,
        shape: vec![n, rows, cols],
        stream: Arc::clone(stream),
    })
}

/// Unsqueeze: add a dimension of size 1 at the given position.
pub fn gpu_unsqueeze(x: &StreamTensor, dim: usize) -> Result<StreamTensor, MLError> {
    if dim > x.shape.len() {
        return Err(MLError::InvalidInput(format!(
            "gpu_unsqueeze dim {} out of range for shape {:?}", dim, x.shape
        )));
    }
    let mut new_shape = x.shape.clone();
    new_shape.insert(dim, 1);
    x.reshape(&new_shape)
}

/// Squeeze: remove a dimension of size 1 at the given position.
pub fn gpu_squeeze(x: &StreamTensor, dim: usize) -> Result<StreamTensor, MLError> {
    if dim >= x.shape.len() {
        return Err(MLError::InvalidInput(format!(
            "gpu_squeeze dim {} out of range for shape {:?}", dim, x.shape
        )));
    }
    if x.dim(dim)? != 1 {
        return Err(MLError::InvalidInput(format!(
            "gpu_squeeze dim {} is not 1 (is {})", dim, x.dim(dim)?
        )));
    }
    let mut new_shape = x.shape.clone();
    new_shape.remove(dim);
    x.reshape(&new_shape)
}

/// Element-wise maximum of tensor and scalar.
pub fn gpu_max_scalar(x: &StreamTensor, scalar: f32) -> Result<StreamTensor, MLError> {
    // clamp(min=scalar, max=f32::MAX) effectively computes max(x, scalar)
    let k = ew(&x.stream)?;
    let out = k.unary(&x.data, x.numel(), 4, scalar, f32::MAX)?; // 4 = clamp
    Ok(wrap(out, &x.shape, &x.stream))
}

/// Element-wise division: a / b.
pub fn gpu_div(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape != b.shape {
        return Err(MLError::DimensionMismatch { expected: a.numel(), actual: b.numel() });
    }
    let k = ew(&a.stream)?;
    let out = k.binary(&a.data, &b.data, a.numel(), 3)?; // 3 = div
    Ok(wrap(out, &a.shape, &a.stream))
}

/// Broadcast add: a (rows, cols) + b (rows, 1) -> (rows, cols).
pub fn gpu_broadcast_add_col(a: &StreamTensor, b: &StreamTensor) -> Result<StreamTensor, MLError> {
    if a.shape.len() != 2 || b.shape.len() != 2 {
        return Err(MLError::InvalidInput("gpu_broadcast_add_col requires 2D tensors".to_owned()));
    }
    let rows = a.dim(0)?;
    let cols = a.dim(1)?;
    if b.dim(0)? != rows || b.dim(1)? != 1 {
        return Err(MLError::DimensionMismatch { expected: rows, actual: b.dim(0)? });
    }
    let k = ew(&a.stream)?;
    let total = rows * cols;
    // broadcast_col: matrix[M,N] op col[M,1], op=2(add), col_is_lhs=0
    // => out[i] = matrix[i] + col[i/cols]
    let out = k.broadcast_col(&a.data, &b.data, total, cols, 2, 0)?;
    Ok(wrap(out, &[rows, cols], &a.stream))
}

/// Identity matrix (2D, square).
pub fn gpu_eye(n: usize, stream: &Arc<CudaStream>) -> Result<StreamTensor, MLError> {
    let mut data = vec![0.0_f32; n * n];
    for i in 0..n {
        if let Some(slot) = data.get_mut(i * n + i) {
            *slot = 1.0;
        }
    }
    StreamTensor::from_vec(data, &[n, n], stream)
}

/// Concatenate tensors along dimension 0 (batch dimension).
///
/// All tensors must have the same shape except for dimension 0.
/// Result shape: [sum(dim0), dim1, dim2, ...].
///
/// Uses GPU DtoD copies -- no host roundtrip.
pub fn gpu_cat_dim0(tensors: &[StreamTensor], stream: &Arc<CudaStream>) -> Result<StreamTensor, MLError> {
    if tensors.is_empty() {
        return Err(MLError::InvalidInput("gpu_cat_dim0: empty list".to_owned()));
    }
    let first = tensors.first().ok_or_else(|| MLError::InvalidInput("empty".to_owned()))?;
    let tail_shape: Vec<usize> = first.shape.iter().skip(1).copied().collect();
    let tail_elems: usize = tail_shape.iter().product::<usize>().max(1);

    let mut total_elems = 0_usize;
    let mut total_batch = 0_usize;
    for t in tensors {
        let t_tail: Vec<usize> = t.shape.iter().skip(1).copied().collect();
        if t_tail != tail_shape {
            return Err(MLError::DimensionMismatch {
                expected: tail_elems,
                actual: t_tail.iter().product::<usize>().max(1),
            });
        }
        total_elems += t.numel();
        total_batch += t.dim(0)?;
    }

    let mut out = stream.alloc_zeros::<f32>(total_elems)
        .map_err(|e| MLError::DeviceError(format!("cat_dim0 alloc: {e}")))?;

    let mut offset = 0_usize;
    for t in tensors {
        let n = t.numel();
        dtod_copy(stream, &t.data, &mut out, 0, offset, n)?;
        offset += n;
    }

    let mut out_shape = vec![total_batch];
    out_shape.extend_from_slice(&tail_shape);
    Ok(StreamTensor {
        data: out,
        shape: out_shape,
        stream: Arc::clone(stream),
    })
}

/// Stack tensors along a new leading dimension.
///
/// All tensors must have the same shape.
/// Result shape: [N, ...original_shape].
///
/// Uses GPU DtoD copies -- no host roundtrip.
pub fn gpu_stack_tensors(tensors: &[StreamTensor], stream: &Arc<CudaStream>) -> Result<StreamTensor, MLError> {
    if tensors.is_empty() {
        return Err(MLError::InvalidInput("gpu_stack_tensors: empty list".to_owned()));
    }
    let first = tensors.first().ok_or_else(|| MLError::InvalidInput("empty".to_owned()))?;
    let elem_shape = first.shape.clone();
    let elem_count = first.numel();
    let total = tensors.len() * elem_count;

    let mut out = stream.alloc_zeros::<f32>(total)
        .map_err(|e| MLError::DeviceError(format!("stack_tensors alloc: {e}")))?;

    for (i, t) in tensors.iter().enumerate() {
        if t.shape != elem_shape {
            return Err(MLError::DimensionMismatch {
                expected: first.numel(),
                actual: t.numel(),
            });
        }
        dtod_copy(stream, &t.data, &mut out, 0, i * elem_count, elem_count)?;
    }

    let mut out_shape = vec![tensors.len()];
    out_shape.extend_from_slice(&elem_shape);
    Ok(StreamTensor {
        data: out,
        shape: out_shape,
        stream: Arc::clone(stream),
    })
}

/// Ones tensor.
pub fn gpu_ones(shape: &[usize], stream: &Arc<CudaStream>) -> Result<StreamTensor, MLError> {
    gpu_full(shape, 1.0, stream)
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_stream_tensor_zeros() {
        let stream = test_stream();
        let t = StreamTensor::zeros(&[2, 3], &stream).unwrap();
        assert_eq!(t.shape, vec![2, 3]);
        assert_eq!(t.numel(), 6);
        let v = t.to_vec().unwrap(); // test-only readback
        assert!(v.iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_stream_tensor_from_vec() {
        let stream = test_stream();
        let t = StreamTensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &stream).unwrap();
        let v = t.to_vec().unwrap(); // test-only readback
        assert_eq!(v, vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_stream_linear_forward() {
        let stream = test_stream();
        let linear = StreamLinear::new(4, 3, &stream).unwrap();
        let input = StreamTensor::from_vec(vec![1.0; 8], &[2, 4], &stream).unwrap();
        let output = linear.forward(&input).unwrap();
        assert_eq!(output.shape, vec![2, 3]);
    }

    #[test]
    fn test_gpu_matmul_correctness() {
        let stream = test_stream();
        let a = StreamTensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], &stream).unwrap();
        let b = StreamTensor::from_vec(vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0], &[3, 2], &stream).unwrap();
        let c = gpu_matmul(&a, &b).unwrap();
        assert_eq!(c.shape, vec![2, 2]);
        let v = c.to_vec().unwrap(); // test-only readback
        assert!((v[0] - 4.0).abs() < 1e-5);
        assert!((v[1] - 5.0).abs() < 1e-5);
        assert!((v[2] - 10.0).abs() < 1e-5);
        assert!((v[3] - 11.0).abs() < 1e-5);
    }

    #[test]
    fn test_gpu_silu_values() {
        let stream = test_stream();
        let x = StreamTensor::from_vec(vec![0.0, 1.0, -1.0], &[3], &stream).unwrap();
        let y = gpu_silu(&x).unwrap();
        let v = y.to_vec().unwrap(); // test-only readback
        assert!((v[0]).abs() < 1e-5);
        assert!((v[1] - 0.7311).abs() < 0.01);
        assert!((v[2] + 0.2689).abs() < 0.01);
    }

    #[test]
    fn test_gpu_cat_dim1_correctness() {
        let stream = test_stream();
        let a = StreamTensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &stream).unwrap();
        let b = StreamTensor::from_vec(vec![5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[2, 3], &stream).unwrap();
        let c = gpu_cat_dim1(&a, &b).unwrap();
        assert_eq!(c.shape, vec![2, 5]);
        let v = c.to_vec().unwrap(); // test-only readback
        assert_eq!(v, vec![1.0, 2.0, 5.0, 6.0, 7.0, 3.0, 4.0, 8.0, 9.0, 10.0]);
    }
}
