//! GPU-resident tensor wrapping `CudaSlice<f32>` with shape metadata.
//!
//! Replaces `candle_core::Tensor` for the CUDA autograd system. This is a thin
//! wrapper — no autograd graph, no dtype dispatch, no device polymorphism.
//! All data is F32 on a single CUDA device.

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

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

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

use crate::MLError;
use crate::native_types::NativeDType;

// ── Pointer helpers (same pattern as linear.rs) ──────────────────────────

/// 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
}

/// Helper trait for scalar values that can be converted to f32.
pub trait ScalarValue {
    fn to_f32_scalar(self) -> f32;
}

impl ScalarValue for f32 {
    fn to_f32_scalar(self) -> f32 { self }
}

impl ScalarValue for f64 {
    fn to_f32_scalar(self) -> f32 { self as f32 }
}

impl ScalarValue for &[f32] {
    fn to_f32_scalar(self) -> f32 { self.first().copied().unwrap_or(0.0) }
}

/// A GPU-resident f32 tensor with shape information.
///
/// This is a lightweight wrapper around `CudaSlice<f32>` that carries shape
/// metadata. Unlike Candle's `Tensor`, it has no autograd graph, no reference
/// counting of operations, and no device abstraction — it is always on GPU.
///
/// ## Ownership
///
/// The inner `CudaSlice<f32>` is reference-counted (cudarc 0.19). Cloning a
/// `GpuTensor` shares the underlying GPU allocation — no device-side copy.
/// The GPU memory is freed when the last reference is dropped.
#[derive(Clone)]
pub struct GpuTensor {
    /// Raw GPU buffer. Element count = product of `shape`.
    pub(crate) data: CudaSlice<f32>,
    /// Tensor dimensions, e.g. `[batch, features]` or `[out_dim, in_dim]`.
    pub(crate) shape: Vec<usize>,
}

impl GpuTensor {
    /// Create a new `GpuTensor` from a pre-allocated `CudaSlice` and shape.
    ///
    /// # Errors
    ///
    /// Returns an error if the product of `shape` does not match `data.len()`.
    pub fn new(data: CudaSlice<f32>, shape: Vec<usize>) -> Result<Self, MLError> {
        let expected: usize = shape.iter().product();
        let actual = data.len();
        if expected != actual {
            return Err(MLError::DimensionMismatch { expected, actual });
        }
        Ok(Self { data, shape })
    }

    /// Access the underlying `CudaSlice<f32>` for direct CUDA kernel operations.
    ///
    /// Prefer higher-level methods (add, sub, affine, argmax, etc.) when available.
    /// Use this only when interfacing with raw CUDA kernels (elementwise, reductions).
    pub fn cuda_data(&self) -> &CudaSlice<f32> {
        &self.data
    }

    /// Wrap an existing `CudaSlice` as a 1-D tensor (no shape validation needed).
    pub fn from_flat(data: CudaSlice<f32>) -> Self {
        let len = data.len();
        Self {
            data,
            shape: vec![len],
        }
    }

    /// Allocate a zero-filled tensor of the given shape on `stream`.
    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::ModelError(format!("GpuTensor::zeros({shape:?}): {e}"))
        })?;
        Ok(Self {
            data,
            shape: shape.to_vec(),
        })
    }

    /// Upload host data to a new GPU tensor.
    pub fn from_host(
        host: &[f32],
        shape: Vec<usize>,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let expected: usize = shape.iter().product();
        if expected != host.len() {
            return Err(MLError::DimensionMismatch {
                expected,
                actual: host.len(),
            });
        }
        let mut data = stream.alloc_zeros::<f32>(expected).map_err(|e| {
            MLError::ModelError(format!("GpuTensor::from_host alloc: {e}"))
        })?;
        stream.memcpy_htod(host, &mut data).map_err(|e| {
            MLError::ModelError(format!("GpuTensor::from_host HtoD: {e}"))
        })?;
        Ok(Self { data, shape })
    }

    /// Download tensor contents to host memory.
    #[cold]
    pub fn to_host(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
        let mut host = vec![0.0_f32; self.numel()];
        stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| { // gpu-exit: API boundary readback
            MLError::ModelError(format!("GpuTensor::to_host DtoH: {e}"))
        })?;
        Ok(host)
    }

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

    /// Tensor shape as a slice.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Number of dimensions.
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    /// Borrow the underlying `CudaSlice`.
    pub fn data(&self) -> &CudaSlice<f32> {
        &self.data
    }

    /// Mutably borrow the underlying `CudaSlice`.
    pub fn data_mut(&mut self) -> &mut CudaSlice<f32> {
        &mut self.data
    }

    /// Consume the tensor and return the underlying `CudaSlice` and shape.
    ///
    /// Useful for converting back to other GPU buffer types (e.g. `CudaVec`)
    /// without a GPU-side copy.
    pub fn into_parts(self) -> (CudaSlice<f32>, Vec<usize>) {
        (self.data, self.shape)
    }

    // -----------------------------------------------------------------------
    // Compatibility shims — replaces candle_core::Tensor factory methods.
    // These are named without the `_candle` suffix as canonical APIs.
    // Legacy `_candle` suffix aliases are provided below for migration.
    // -----------------------------------------------------------------------

    /// Create a tensor filled with random normal values (mean=0, std=`std_dev`).
    pub fn randn(shape: &[usize], std_dev: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n: usize = shape.iter().product();
        let mut host = vec![0.0_f32; n];
        // Use thread_rng for initialization (host-side, not hot path)
        let mut rng = rand::thread_rng();
        let dist = rand_distr::Normal::new(0.0_f32, std_dev).map_err(|e| {
            MLError::ModelError(format!("randn: invalid std_dev={std_dev}: {e}"))
        })?;
        for v in &mut host {
            use rand_distr::Distribution;
            *v = dist.sample(&mut rng);
        }
        Self::from_host(&host, shape.to_vec(), stream)
    }

    /// Create a scalar tensor (single-element).
    pub fn scalar(value: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Self::from_host(&[value], vec![1], stream)
    }

    /// Create an identity matrix of size `n x n`.
    pub fn eye(n: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let mut host = vec![0.0_f32; n * n];
        for i in 0..n {
            if let Some(elem) = host.get_mut(i * n + i) {
                *elem = 1.0;
            }
        }
        Self::from_host(&host, vec![n, n], stream)
    }

    /// Concatenate tensors along a given dimension.
    ///
    /// For now this downloads to host, concatenates, and re-uploads.
    /// TODO: GPU-side concat kernel for hot path.
    pub fn cat(tensors: &[&Self], dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if tensors.is_empty() {
            return Err(MLError::ModelError("cat: empty tensor list".to_string()));
        }
        let ndim = tensors.first().map_or(0, |t| t.ndim());
        if dim >= ndim {
            return Err(MLError::ModelError(format!(
                "cat: dim {dim} >= ndim {ndim}"
            )));
        }

        let mut total_dim = 0_usize;
        for t in tensors {
            total_dim += t.shape().get(dim).copied().unwrap_or(0);
        }

        let ref_shape = tensors.first().map_or(&[] as &[usize], |t| t.shape());
        let mut new_shape = ref_shape.to_vec();
        new_shape[dim] = total_dim;
        let total_elems: usize = new_shape.iter().product();

        // GPU-native: allocate output and DtoD copy each tensor's data
        let mut output = stream.alloc_zeros::<f32>(total_elems).map_err(|e| {
            MLError::ModelError(format!("cat: alloc {total_elems} f32: {e}"))
        })?;

        if dim == 0 {
            // dim=0: tensors are contiguous blocks — sequential DtoD copies
            let mut offset = 0_usize;
            for t in tensors {
                let n = t.numel();
                if n > 0 {
                    let src_view = t.data.slice(..n);
                    let mut dst_view = output.slice_mut(offset..offset + n);
                    stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| {
                        MLError::ModelError(format!("cat dim=0 DtoD: {e}"))
                    })?;
                }
                offset += n;
            }
        } else {
            // dim>0: interleaved copy via per-row DtoD
            let inner: usize = if dim + 1 < ref_shape.len() {
                ref_shape[dim + 1..].iter().product()
            } else {
                1
            };
            let outer: usize = ref_shape[..dim].iter().product();
            let dst_stride = total_dim * inner;

            for o in 0..outer {
                let mut col_offset = 0_usize;
                for t in tensors.iter() {
                    let d = t.shape().get(dim).copied().unwrap_or(0);
                    let src_stride = d * inner;
                    let src_start = o * src_stride;
                    let dst_start = o * dst_stride + col_offset;
                    let n = src_stride;
                    if n > 0 && src_start + n <= t.data.len() && dst_start + n <= output.len() {
                        let src_view = t.data.slice(src_start..src_start + n);
                        let mut dst_view = output.slice_mut(dst_start..dst_start + n);
                        stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| {
                            MLError::ModelError(format!("cat dim={dim} DtoD: {e}"))
                        })?;
                    }
                    col_offset += src_stride;
                }
            }
        }

        Self::new(output, new_shape)
    }

    /// Stack tensors along a new dimension 0 (GPU-native DtoD).
    pub fn stack(tensors: &[&Self], _dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if tensors.is_empty() {
            return Err(MLError::ModelError("stack: empty tensor list".to_string()));
        }
        let inner_shape = tensors.first().map_or(&[][..], |t| t.shape());
        let batch = tensors.len();
        let elems_per: usize = inner_shape.iter().product();
        let total = batch * elems_per;

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

        for (i, t) in tensors.iter().enumerate() {
            let n = t.numel();
            if n > 0 {
                let src_view = t.data.slice(..n);
                let mut dst_view = output.slice_mut(i * elems_per..(i + 1) * elems_per);
                stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| {
                    MLError::ModelError(format!("stack DtoD [{i}]: {e}"))
                })?;
            }
        }

        let mut new_shape = vec![batch];
        new_shape.extend_from_slice(inner_shape);
        Self::new(output, new_shape)
    }

    // -----------------------------------------------------------------------
    // _candle suffix aliases removed — migration complete, zero callers remain.

    // -----------------------------------------------------------------------
    // Tensor algebra: element-wise ops, reductions, slicing.
    // All ops are GPU-native — no host downloads except to_host/to_scalar.
    // -----------------------------------------------------------------------

    /// Clone this tensor by copying GPU data to a new allocation.
    ///
    /// Uses async device-to-device memcpy — zero host download.
    pub fn gpu_clone(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n = self.numel();
        let dst = stream.alloc_zeros::<f32>(n).map_err(|e| {
            MLError::ModelError(format!("gpu_clone alloc: {e}"))
        })?;
        let num_bytes = n * std::mem::size_of::<f32>();
        let src_ptr = raw_ptr(&self.data, stream);
        let dst_ptr = raw_ptr(&dst, stream);
        // SAFETY: src and dst are valid device allocations on the same context.
        // num_bytes = n * sizeof(f32) does not exceed either allocation.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| MLError::ModelError(format!("gpu_clone DtoD: {e}")))?;
        }
        Ok(Self {
            data: dst,
            shape: self.shape.clone(),
        })
    }

    /// No-op dtype cast — `GpuTensor` is always F32. Returns a clone.
    pub fn to_dtype(&self, _dtype: NativeDType, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        self.gpu_clone(stream)
    }

    /// No-op detach (with stream) — no autograd tape. Returns a deep clone.
    pub fn detach_with_stream(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        self.gpu_clone(stream)
    }

    /// No-op detach — no autograd tape. Returns a shallow clone (shared GPU memory).
    ///
    /// This is the Candle-compatible signature that does not require a stream
    /// argument. Since `GpuTensor` has no autograd graph, detach is semantically
    /// a no-op and a cheap ref-counted clone is safe.
    pub fn detach(&self) -> Self {
        self.clone()
    }

    /// Create a zero-filled tensor with the same shape as `self`.
    pub fn zeros_like(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        Self::zeros(&self.shape, stream)
    }

    /// Create a tensor filled with a constant value.
    pub fn full(shape: &[usize], value: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n: usize = shape.iter().product();
        let host = vec![value; n];
        Self::from_host(&host, shape.to_vec(), stream)
    }

    /// Create a 1-D tensor with values in `[start, end)` stepping by `step`.
    pub fn arange(start: f32, end: f32, step: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if step == 0.0 {
            return Err(MLError::InvalidInput("arange: step cannot be zero".to_string()));
        }
        let n = ((end - start) / step).ceil().max(0.0) as usize;
        let mut host = Vec::with_capacity(n);
        let mut v = start;
        for _ in 0..n {
            host.push(v);
            v += step;
        }
        let len = host.len();
        Self::from_host(&host, vec![len], stream)
    }

    /// Element-wise addition. Shapes must match exactly.
    ///
    /// GPU-native: uses `elementwise_binary` kernel (op=0).
    pub fn add(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if self.shape != other.shape {
            return Err(MLError::DimensionMismatch {
                expected: self.numel(),
                actual: other.numel(),
            });
        }
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.binary(&self.data, &other.data, self.numel(), 0)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise subtraction. Shapes must match exactly.
    ///
    /// GPU-native: uses `elementwise_binary` kernel (op=1).
    pub fn sub(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if self.shape != other.shape {
            return Err(MLError::DimensionMismatch {
                expected: self.numel(),
                actual: other.numel(),
            });
        }
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.binary(&self.data, &other.data, self.numel(), 1)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise multiply. Shapes must match exactly.
    ///
    /// GPU-native: uses `elementwise_binary` kernel (op=2).
    pub fn mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if self.shape != other.shape {
            return Err(MLError::DimensionMismatch {
                expected: self.numel(),
                actual: other.numel(),
            });
        }
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.binary(&self.data, &other.data, self.numel(), 2)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise multiply with broadcasting.
    ///
    /// Supports the common cases:
    /// - Same shape: plain element-wise multiply.
    /// - Scalar (numel=1) * tensor: broadcast the scalar.
    /// - `[1, N]` * `[M, N]` or `[M, N]` * `[1, N]`: broadcast rows.
    ///
    /// GPU-native: all paths use CUDA kernels with zero host download.
    pub fn broadcast_mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        // Fast path: identical shapes
        if self.shape == other.shape {
            return self.mul(other, stream);
        }
        let kernels = super::elementwise::get_or_compile(stream)?;
        // Scalar broadcast
        if other.numel() == 1 {
            let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 0, 0)?;
            return Ok(Self { data: out, shape: self.shape.clone() });
        }
        if self.numel() == 1 {
            let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 0, 1)?;
            return Ok(Self { data: out, shape: other.shape.clone() });
        }
        // Row broadcast: [1, N] * [M, N] or [M, N] * [1, N]
        // Column broadcast: [M, 1] * [M, N] or [M, N] * [M, 1]
        if self.ndim() == 2 && other.ndim() == 2 {
            let (m_a, n_a) = (
                self.shape.first().copied().unwrap_or(0),
                self.shape.get(1).copied().unwrap_or(0),
            );
            let (m_b, n_b) = (
                other.shape.first().copied().unwrap_or(0),
                other.shape.get(1).copied().unwrap_or(0),
            );
            if n_a == n_b {
                if m_a == 1 {
                    // [1, N] * [M, N]: self is the row, other is the matrix
                    let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 0, 0)?;
                    return Ok(Self { data: out, shape: other.shape.clone() });
                }
                if m_b == 1 {
                    // [M, N] * [1, N]: other is the row, self is the matrix
                    let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 0, 0)?;
                    return Ok(Self { data: out, shape: self.shape.clone() });
                }
            }
            if m_a == m_b {
                if n_b == 1 {
                    // [M, N] * [M, 1]: other is the column vector, self is the matrix
                    let out = kernels.broadcast_col(&self.data, &other.data, m_a * n_a, n_a, 0, 0)?;
                    return Ok(Self { data: out, shape: self.shape.clone() });
                }
                if n_a == 1 {
                    // [M, 1] * [M, N]: self is the column vector, other is the matrix
                    let out = kernels.broadcast_col(&other.data, &self.data, m_b * n_b, n_b, 0, 0)?;
                    return Ok(Self { data: out, shape: other.shape.clone() });
                }
            }
        }
        Err(MLError::ModelError(format!(
            "broadcast_mul: incompatible shapes {:?} vs {:?}",
            self.shape, other.shape
        )))
    }

    /// Element-wise divide with broadcasting.
    ///
    /// Supports the same broadcast patterns as `broadcast_mul`.
    ///
    /// GPU-native: all paths use CUDA kernels with zero host download.
    pub fn broadcast_div(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        // Fast path: identical shapes
        if self.shape == other.shape {
            let out = kernels.binary(&self.data, &other.data, self.numel(), 3)?;
            return Ok(Self { data: out, shape: self.shape.clone() });
        }
        // Scalar broadcast (divisor is scalar)
        if other.numel() == 1 {
            // tensor / scalar: scalar_is_lhs=0, op=1 (div)
            let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 1, 0)?;
            return Ok(Self { data: out, shape: self.shape.clone() });
        }
        if self.numel() == 1 {
            // scalar / tensor: scalar_is_lhs=1, op=1 (div)
            let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 1, 1)?;
            return Ok(Self { data: out, shape: other.shape.clone() });
        }
        // Row broadcast for 2-D
        // Column broadcast for 2-D
        if self.ndim() == 2 && other.ndim() == 2 {
            let (m_a, n_a) = (
                self.shape.first().copied().unwrap_or(0),
                self.shape.get(1).copied().unwrap_or(0),
            );
            let (m_b, n_b) = (
                other.shape.first().copied().unwrap_or(0),
                other.shape.get(1).copied().unwrap_or(0),
            );
            if n_a == n_b {
                if m_b == 1 {
                    // [M, N] / [1, N]: matrix=self, row=other, op=1 (div), row_is_lhs=0
                    let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 1, 0)?;
                    return Ok(Self { data: out, shape: self.shape.clone() });
                }
                if m_a == 1 {
                    // [1, N] / [M, N]: matrix=other, row=self, op=1 (div), row_is_lhs=1
                    let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 1, 1)?;
                    return Ok(Self { data: out, shape: other.shape.clone() });
                }
            }
            if m_a == m_b {
                if n_b == 1 {
                    // [M, N] / [M, 1]: matrix=self, col_vec=other, op=1 (div), col_is_lhs=0
                    let out = kernels.broadcast_col(&self.data, &other.data, m_a * n_a, n_a, 1, 0)?;
                    return Ok(Self { data: out, shape: self.shape.clone() });
                }
                if n_a == 1 {
                    // [M, 1] / [M, N]: matrix=other, col_vec=self, op=1 (div), col_is_lhs=1
                    let out = kernels.broadcast_col(&other.data, &self.data, m_b * n_b, n_b, 1, 1)?;
                    return Ok(Self { data: out, shape: other.shape.clone() });
                }
            }
        }
        Err(MLError::ModelError(format!(
            "broadcast_div: incompatible shapes {:?} vs {:?}",
            self.shape, other.shape
        )))
    }

    /// Matrix multiplication via cuBLAS sgemm.
    ///
    /// For 2-D tensors `[M, K]` @ `[K, N]` -> `[M, N]`.
    /// For 1-D `[K]` @ `[K, N]` -> `[N]` (treated as `[1, K]` matmul, then squeeze).
    pub fn matmul(
        &self,
        other: &Self,
        cublas: &CudaBlas,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let (m, k_a) = match self.ndim() {
            1 => (1_usize, self.shape.first().copied().unwrap_or(0)),
            2 => (
                self.shape.first().copied().unwrap_or(0),
                self.shape.get(1).copied().unwrap_or(0),
            ),
            _ => {
                return Err(MLError::ModelError(format!(
                    "matmul: lhs must be 1-D or 2-D, got {:?}",
                    self.shape
                )));
            }
        };
        let (k_b, n) = match other.ndim() {
            2 => (
                other.shape.first().copied().unwrap_or(0),
                other.shape.get(1).copied().unwrap_or(0),
            ),
            _ => {
                return Err(MLError::ModelError(format!(
                    "matmul: rhs must be 2-D, got {:?}",
                    other.shape
                )));
            }
        };
        if k_a != k_b {
            return Err(MLError::DimensionMismatch {
                expected: k_a,
                actual: k_b,
            });
        }

        let mut c = Self::zeros(&[m, n], stream)?;

        // cuBLAS column-major: C_col[N, M] = B_col[N, K] @ A_col[K, M]
        // Row-major A[M,K] stored as col-major A_col[K,M] (transposed)
        // So: transA=N (A_col is already [K,M]), transB=N for B_col[N,K]
        // Wait — we have row-major data but cuBLAS expects col-major.
        // Row-major C[M,N] = A[M,K] @ B[K,N]
        // In col-major view: C_col[N,M] = B_col[N,K] @ A_col[K,M]
        //   transA=N, transB=N, m=N, n=M, k=K, lda=N, ldb=K, ldc=N
        let a_ptr = raw_ptr(&self.data, stream);
        let b_ptr = raw_ptr(&other.data, stream);
        let c_ptr = raw_ptr_mut(&mut c.data, stream);

        unsafe {
            cudarc::cublas::result::sgemm(
                *cublas.handle(),
                cublasOperation_t::CUBLAS_OP_N, // transA (B in row-major)
                cublasOperation_t::CUBLAS_OP_N, // transB (A in row-major)
                n as i32,                       // m (cols of C)
                m as i32,                       // n (rows of C)
                k_a as i32,                     // k
                &1.0_f32 as *const f32,
                b_ptr as *const f32,            // A in cuBLAS = B in row-major
                n as i32,                       // lda = N
                a_ptr as *const f32,            // B in cuBLAS = A in row-major
                k_a as i32,                     // ldb = K
                &0.0_f32 as *const f32,
                c_ptr as *mut f32,
                n as i32,                       // ldc = N
            )
            .map_err(|e| MLError::ModelError(format!("cuBLAS sgemm matmul: {e:?}")))?;
        }

        // If lhs was 1-D, squeeze the leading dim
        if self.ndim() == 1 {
            c = c.reshape(vec![n])?;
        }
        Ok(c)
    }

    /// Add a dimension of size 1 at position `dim`.
    pub fn unsqueeze(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if dim > self.ndim() {
            return Err(MLError::ModelError(format!(
                "unsqueeze: dim {dim} > ndim {}",
                self.ndim()
            )));
        }
        let mut new_shape = self.shape.clone();
        new_shape.insert(dim, 1);
        // Data is the same, just reshape
        let cloned = self.gpu_clone(stream)?;
        Ok(Self {
            data: cloned.data,
            shape: new_shape,
        })
    }

    /// Remove a dimension of size 1 at position `dim`.
    pub fn squeeze(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if dim >= self.ndim() {
            return Err(MLError::ModelError(format!(
                "squeeze: dim {dim} >= ndim {}",
                self.ndim()
            )));
        }
        if self.shape.get(dim).copied().unwrap_or(0) != 1 {
            return Err(MLError::ModelError(format!(
                "squeeze: dim {dim} has size {}, expected 1",
                self.shape.get(dim).copied().unwrap_or(0)
            )));
        }
        let mut new_shape = self.shape.clone();
        new_shape.remove(dim);
        if new_shape.is_empty() {
            new_shape.push(1);
        }
        let cloned = self.gpu_clone(stream)?;
        Ok(Self {
            data: cloned.data,
            shape: new_shape,
        })
    }

    /// Slice along a dimension: extract `[start..start+len)` along `dim`.
    ///
    /// For dim=0 (contiguous), uses async DtoD memcpy — zero host download.
    /// For other dims, uses the gather kernel on GPU.
    pub fn narrow(
        &self,
        dim: usize,
        start: usize,
        len: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        if dim >= self.ndim() {
            return Err(MLError::ModelError(format!(
                "narrow: dim {dim} >= ndim {}",
                self.ndim()
            )));
        }
        let dim_size = self.shape.get(dim).copied().unwrap_or(0);
        if start + len > dim_size {
            return Err(MLError::ModelError(format!(
                "narrow: start({start}) + len({len}) > dim_size({dim_size})"
            )));
        }

        // Compute strides for row-major layout
        let ndim = self.ndim();
        let mut strides = vec![1_usize; ndim];
        for i in (0..ndim.saturating_sub(1)).rev() {
            let next = strides.get(i + 1).copied().unwrap_or(1);
            let sh = self.shape.get(i + 1).copied().unwrap_or(1);
            if let Some(s) = strides.get_mut(i) {
                *s = next * sh;
            }
        }

        let mut new_shape = self.shape.clone();
        if let Some(d) = new_shape.get_mut(dim) {
            *d = len;
        }
        let new_numel: usize = new_shape.iter().product();

        let stride = strides.get(dim).copied().unwrap_or(1);

        // For dim=0 this is a contiguous sub-range — use DtoD memcpy
        if dim == 0 {
            let offset = start * stride;
            let src_view = self.data.slice(offset..offset + new_numel);
            let dst = stream.alloc_zeros::<f32>(new_numel).map_err(|e| {
                MLError::ModelError(format!("narrow alloc: {e}"))
            })?;
            let num_bytes = new_numel * std::mem::size_of::<f32>();
            let dst_ptr = raw_ptr(&dst, stream);
            let (src_ptr, _src_sync) = src_view.device_ptr(stream);
            // SAFETY: src_view and dst are valid device allocations on the same context.
            // num_bytes does not exceed either allocation.
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("narrow DtoD: {e}")))?;
            }
            return Ok(Self { data: dst, shape: new_shape });
        }

        // General case: use gather kernel with sequential indices
        let indices: Vec<u32> = (start..start + len).map(|i| i as u32).collect();
        let mut indices_dev = stream.alloc_zeros::<u32>(len).map_err(|e| {
            MLError::ModelError(format!("narrow indices alloc: {e}"))
        })?;
        stream.memcpy_htod(&indices, &mut indices_dev).map_err(|e| {
            MLError::ModelError(format!("narrow indices HtoD: {e}"))
        })?;

        let inner_size = stride; // elements per index along this dim
        let outer_count = if dim == 0 { 1 } else {
            self.shape.iter().take(dim).product::<usize>()
        };

        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.gather(&self.data, &indices_dev, len, inner_size, outer_count, dim_size)?;
        Ok(Self { data: out, shape: new_shape })
    }

    /// Argmax along a dimension. Returns indices as `Vec<u32>`.
    ///
    /// For a 2-D tensor `[M, N]` with `dim=1`, returns `M` indices in `[0..N)`.
    ///
    /// GPU-native: delegates to `ReductionKernels::argmax` / `argmax_rows`.
    pub fn argmax(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Vec<u32>, MLError> {
        if dim >= self.ndim() {
            return Err(MLError::ModelError(format!(
                "argmax: dim {dim} >= ndim {}",
                self.ndim()
            )));
        }
        let reductions = super::reductions::ReductionKernels::new(stream)?;

        // For 1-D, return single argmax via flat kernel
        if self.ndim() == 1 {
            let idx = reductions.argmax(&self.data, self.numel())?;
            return Ok(vec![idx]);
        }

        // For 2-D with dim=1: argmax per row — use ReductionKernels::argmax_rows
        if self.ndim() == 2 && dim == 1 {
            let rows = self.shape.first().copied().unwrap_or(0);
            let cols = self.shape.get(1).copied().unwrap_or(0);
            return reductions.argmax_rows(&self.data, rows, cols);
        }

        // For 2-D with dim=0: transpose then argmax per row of transposed
        if self.ndim() == 2 && dim == 0 {
            let rows = self.shape.first().copied().unwrap_or(0);
            let cols = self.shape.get(1).copied().unwrap_or(0);
            let kernels = super::elementwise::get_or_compile(stream)?;
            let transposed = kernels.transpose_2d(&self.data, rows, cols)?;
            // Now transposed is [cols, rows]. argmax per row of [cols, rows] gives
            // per-original-column argmax along original rows (dim=0).
            return reductions.argmax_rows(&transposed, cols, rows);
        }

        Err(MLError::ModelError(format!(
            "argmax: unsupported ndim={} dim={dim}",
            self.ndim()
        )))
    }

    /// Extract a single scalar from a 1-element tensor.
    pub fn to_scalar(&self, stream: &Arc<CudaStream>) -> Result<f32, MLError> {
        if self.numel() != 1 {
            return Err(MLError::ModelError(format!(
                "to_scalar: expected 1 element, got {}",
                self.numel()
            )));
        }
        let host = self.to_host(stream)?;
        Ok(host.first().copied().unwrap_or(0.0))
    }

    /// Mean of all elements.
    ///
    /// GPU-native: delegates to `ReductionKernels::stats().mean`.
    pub fn mean_all(&self, stream: &Arc<CudaStream>) -> Result<f32, MLError> {
        let n = self.numel();
        if n == 0 {
            return Err(MLError::ModelError("mean_all: empty tensor".to_string()));
        }
        let reductions = super::reductions::ReductionKernels::new(stream)?;
        let stats = reductions.stats(&self.data, n)?;
        Ok(stats.mean)
    }

    // -----------------------------------------------------------------------
    // Compatibility shims — methods expected by the `ml` crate callers.
    // -----------------------------------------------------------------------

    /// Alias for `shape()` — some callers use `dims()` (Candle convention).
    pub fn dims(&self) -> &[usize] {
        self.shape()
    }

    /// Return size of a single dimension by index.
    ///
    /// # Errors
    ///
    /// Returns an error if `dim` is out of range.
    pub fn dim(&self, dim: usize) -> Result<usize, MLError> {
        self.shape.get(dim).copied().ok_or_else(|| {
            MLError::DimensionMismatch {
                expected: self.ndim(),
                actual: dim,
            }
        })
    }

    /// Return `(d0, d1)` for a 2-D tensor.
    pub fn dims2(&self) -> Result<(usize, usize), MLError> {
        if self.ndim() != 2 {
            return Err(MLError::ModelError(format!(
                "dims2: expected 2-D tensor, got {:?}", self.shape
            )));
        }
        Ok((
            self.shape.first().copied().unwrap_or(0),
            self.shape.get(1).copied().unwrap_or(0),
        ))
    }

    /// Return `(d0, d1, d2)` for a 3-D tensor.
    pub fn dims3(&self) -> Result<(usize, usize, usize), MLError> {
        if self.ndim() != 3 {
            return Err(MLError::ModelError(format!(
                "dims3: expected 3-D tensor, got {:?}", self.shape
            )));
        }
        Ok((
            self.shape.first().copied().unwrap_or(0),
            self.shape.get(1).copied().unwrap_or(0),
            self.shape.get(2).copied().unwrap_or(0),
        ))
    }

    /// Create a GPU tensor from a `Vec<f32>`, uploading to device.
    ///
    /// Alias for `from_host()` — callers migrating from Candle used `from_vec()`.
    pub fn from_vec<S: Into<Vec<usize>>>(
        data: Vec<f32>,
        shape: S,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let shape = shape.into();
        Self::from_host(&data, shape, stream)
    }

    /// Create a GPU tensor from a slice, uploading to device.
    ///
    /// Alias for `from_host()` — callers migrating from Candle used `from_slice()`.
    pub fn from_slice<S: Into<Vec<usize>>>(
        data: &[f32],
        shape: S,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        Self::from_host(data, shape.into(), stream)
    }

    /// No-op detach that does NOT require a stream argument.
    ///
    /// Since `GpuTensor` has no autograd tape, detach is semantically a no-op.
    /// This variant avoids cloning and returns a new tensor wrapping the same
    /// logical data (actually a cheap host-roundtrip clone for safety).
    ///
    /// Callers that already have a stream should prefer `detach_with_stream()`
    /// for clarity, but this version exists for Candle API compatibility.
    pub fn detach_no_stream(&self) -> Self {
        // CudaSlice is reference-counted in cudarc 0.19, so cloning shares
        // the underlying GPU allocation without a device-side copy.
        // For a proper deep copy, use gpu_clone(stream).
        Self {
            data: self.data.clone(),
            shape: self.shape.clone(),
        }
    }

    /// Element-wise power.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=0, param1=exp).
    pub fn powf(&self, exp: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.unary(&self.data, self.numel(), 0, exp, 0.0)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise square (`x^2`).
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=1) for `v*v`.
    pub fn sqr(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.unary(&self.data, self.numel(), 1, 0.0, 0.0)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise floor.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=2).
    pub fn floor(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.unary(&self.data, self.numel(), 2, 0.0, 0.0)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Extract a single scalar from a 1-element tensor, converting to type `T`.
    ///
    /// Alias for `to_scalar()` with type conversion. Callers use
    /// `tensor.to_vec0::<f32>(stream)?` as a Candle migration pattern.
    pub fn to_vec0<T: From<f32>>(&self, stream: &Arc<CudaStream>) -> Result<T, MLError> {
        let val = self.to_scalar(stream)?;
        Ok(T::from(val))
    }

    /// Download tensor contents to a `Vec<f32>` (1-D view).
    ///
    /// Alias for `to_host()` — Candle callers used `to_vec1()`.
    pub fn to_vec1(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
        self.to_host(stream)
    }

    /// No-op backward stub. GpuTensor has no autograd tape.
    ///
    /// Returns `Ok(())` unconditionally. Callers that were using Candle's
    /// `loss.backward()` can call this without error while the manual backward
    /// pass is wired up.
    pub fn backward(&self) -> Result<(), MLError> {
        Ok(())
    }

    /// Element-wise ReLU: `max(0, x)`.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=3).
    pub fn relu(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.unary(&self.data, self.numel(), 3, 0.0, 0.0)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Sum of all elements, returned as a scalar tensor.
    ///
    /// GPU-native: delegates to `ReductionKernels::sum()`.
    pub fn sum_all(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n = self.numel();
        if n == 0 {
            return Self::scalar(0.0, stream);
        }
        let reductions = super::reductions::ReductionKernels::new(stream)?;
        let total = reductions.sum(&self.data, n)?;
        Self::scalar(total, stream)
    }

    /// Flatten all dimensions into a single 1-D tensor.
    pub fn flatten_all(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let n = self.numel();
        self.gpu_clone(stream).and_then(|t| t.reshape(vec![n]))
    }

    /// Transpose a 2-D tensor (swap dims 0 and 1).
    ///
    /// GPU-native: uses `transpose_2d` CUDA kernel.
    pub fn transpose(&self, dim0: usize, dim1: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if self.ndim() != 2 || dim0 > 1 || dim1 > 1 || dim0 == dim1 {
            return Err(MLError::ModelError(format!(
                "transpose: only 2-D dim swap supported, got shape {:?} dims ({dim0},{dim1})",
                self.shape
            )));
        }
        let rows = self.shape.first().copied().unwrap_or(0);
        let cols = self.shape.get(1).copied().unwrap_or(0);
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.transpose_2d(&self.data, rows, cols)?;
        Ok(Self { data: out, shape: vec![cols, rows] })
    }

    /// Element-wise absolute value (`|x|`).
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=5).
    pub fn abs(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.abs(&self.data, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise negation (`-x`).
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=6).
    pub fn neg(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.neg(&self.data, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise exponential (`e^x`).
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=7).
    pub fn exp(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.exp(&self.data, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise natural logarithm (`ln(x)`).
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=8).
    pub fn log(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.log(&self.data, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise symlog: `sign(x) * ln(|x| + 1)`.
    ///
    /// GPU-native: computed via abs, affine(+1), log, div, mul kernels.
    /// Zero CPU roundtrip -- all arithmetic stays on GPU.
    pub fn symlog(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let n = self.numel();
        // |x|
        let abs_x = kernels.abs(&self.data, n)?;
        // |x| + 1
        let abs_plus_1 = kernels.affine(&abs_x, 1.0, 1.0, n)?;
        // ln(|x| + 1)
        let ln_val = kernels.log(&abs_plus_1, n)?;
        // |x| + eps  (avoid div-by-zero for sign computation)
        let abs_plus_eps = kernels.affine(&abs_x, 1.0, 1e-8, n)?;
        // sign(x) = x / (|x| + eps)
        let sign_x = kernels.binary(&self.data, &abs_plus_eps, n, 3)?; // op 3 = div
        // sign(x) * ln(|x| + 1)
        let result = kernels.binary(&sign_x, &ln_val, n, 2)?; // op 2 = mul
        Ok(Self { data: result, shape: self.shape.clone() })
    }

    /// Element-wise less-or-equal comparison. Returns 1.0 where `x <= threshold`, else 0.0.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=9).
    pub fn le(&self, threshold: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.le(&self.data, threshold, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise affine transform: `x * mul + add`.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=10).
    pub fn affine(&self, mul: f64, add: f64, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.affine(&self.data, mul as f32, add as f32, self.numel())?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Element-wise subtraction with broadcasting.
    ///
    /// Supports the common cases:
    /// - Same shape: plain element-wise sub.
    /// - Scalar (numel=1): broadcast the scalar.
    /// - `[M, N]` - `[1, N]`: row broadcast.
    ///
    /// GPU-native: all paths use CUDA kernels with zero host download.
    pub fn broadcast_sub(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        // Fast path: identical shapes
        if self.shape == other.shape {
            return self.sub(other, stream);
        }
        let kernels = super::elementwise::get_or_compile(stream)?;
        // Scalar broadcast: self - scalar
        if other.numel() == 1 {
            // tensor - scalar: op=3 (sub), scalar_is_lhs=0
            let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 3, 0)?;
            return Ok(Self { data: out, shape: self.shape.clone() });
        }
        // Scalar broadcast: scalar - tensor
        if self.numel() == 1 {
            // scalar - tensor: op=3 (sub), scalar_is_lhs=1
            let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 3, 1)?;
            return Ok(Self { data: out, shape: other.shape.clone() });
        }
        // Same numel: just use element-wise sub
        if self.numel() == other.numel() {
            let out = kernels.binary(&self.data, &other.data, self.numel(), 1)?;
            return Ok(Self { data: out, shape: self.shape.clone() });
        }
        Err(MLError::ModelError(format!(
            "broadcast_sub: incompatible shapes {:?} vs {:?}",
            self.shape, other.shape
        )))
    }

    /// Element-wise clamp to `[min, max]`.
    ///
    /// GPU-native: uses `elementwise_unary` kernel (op=4, param1=min, param2=max).
    pub fn clamp(&self, min: f32, max: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.unary(&self.data, self.numel(), 4, min, max)?;
        Ok(Self { data: out, shape: self.shape.clone() })
    }

    /// Expand (broadcast) this tensor to match a target shape.
    ///
    /// Only supports expanding dimensions of size 1.
    ///
    /// GPU-native: uploads stride tables to GPU, runs `expand_broadcast` kernel.
    pub fn expand(&self, target_shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if self.ndim() != target_shape.len() {
            return Err(MLError::DimensionMismatch {
                expected: target_shape.len(),
                actual: self.ndim(),
            });
        }
        // Check that each dim is either equal or the source dim is 1
        for (i, (&src, &tgt)) in self.shape.iter().zip(target_shape.iter()).enumerate() {
            if src != tgt && src != 1 {
                return Err(MLError::ModelError(format!(
                    "expand: dim {i} is {src}, cannot expand to {tgt} (must be 1)"
                )));
            }
        }

        let ndim = self.ndim();
        let total: usize = target_shape.iter().product();

        // Compute per-element strides for source and target
        let mut src_strides_host = vec![1_u32; ndim];
        let mut tgt_strides_host = vec![1_u32; ndim];
        let mut src_shape_host: Vec<u32> = self.shape.iter().map(|&s| s as u32).collect();
        // Pad to at least ndim elements (always should be correct)
        while src_shape_host.len() < ndim {
            src_shape_host.push(1);
        }

        for i in (0..ndim.saturating_sub(1)).rev() {
            let sh = self.shape.get(i + 1).copied().unwrap_or(1) as u32;
            let th = target_shape.get(i + 1).copied().unwrap_or(1) as u32;
            let next_src = src_strides_host.get(i + 1).copied().unwrap_or(1);
            let next_tgt = tgt_strides_host.get(i + 1).copied().unwrap_or(1);
            if let Some(ss) = src_strides_host.get_mut(i) {
                *ss = next_src * sh;
            }
            if let Some(ts) = tgt_strides_host.get_mut(i) {
                *ts = next_tgt * th;
            }
        }

        // For expand, source strides for dims of size 1 must be 0 so the kernel
        // reads the same element. But our kernel handles this via src_shape check,
        // so we need per-element strides (stride within the source buffer).
        // The kernel computes: src_coord = (src_dim==1) ? 0 : coord
        // Then: src_idx += src_coord * src_strides[d]
        // We need src_strides[d] to be the stride *within the source allocation*,
        // i.e. for a source of shape [1, N], src_strides = [N, 1].
        // This is already what we computed above (product of subsequent dims).

        // Upload stride tables to GPU
        let mut src_strides_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
            MLError::ModelError(format!("expand src_strides alloc: {e}"))
        })?;
        let mut src_shape_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
            MLError::ModelError(format!("expand src_shape alloc: {e}"))
        })?;
        let mut tgt_strides_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
            MLError::ModelError(format!("expand tgt_strides alloc: {e}"))
        })?;
        stream.memcpy_htod(&src_strides_host, &mut src_strides_dev).map_err(|e| {
            MLError::ModelError(format!("expand src_strides HtoD: {e}"))
        })?;
        stream.memcpy_htod(&src_shape_host, &mut src_shape_dev).map_err(|e| {
            MLError::ModelError(format!("expand src_shape HtoD: {e}"))
        })?;
        stream.memcpy_htod(&tgt_strides_host, &mut tgt_strides_dev).map_err(|e| {
            MLError::ModelError(format!("expand tgt_strides HtoD: {e}"))
        })?;

        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.expand(
            &self.data,
            &src_strides_dev,
            &src_shape_dev,
            &tgt_strides_dev,
            total,
            ndim,
        )?;
        Ok(Self { data: out, shape: target_shape.to_vec() })
    }

    /// Alias for `expand()` — Candle callers use `broadcast_as`.
    pub fn broadcast_as(&self, shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        self.expand(shape, stream)
    }

    /// Select specific indices along a dimension.
    ///
    /// GPU-native: uploads indices to GPU, runs `gather_select` kernel.
    pub fn index_select(
        &self,
        dim: usize,
        indices: &[u32],
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        if dim >= self.ndim() {
            return Err(MLError::ModelError(format!(
                "index_select: dim {dim} >= ndim {}", self.ndim()
            )));
        }
        let dim_size = self.shape.get(dim).copied().unwrap_or(0);

        // Compute strides
        let ndim = self.ndim();
        let mut strides = vec![1_usize; ndim];
        for i in (0..ndim.saturating_sub(1)).rev() {
            let next_stride = strides.get(i + 1).copied().unwrap_or(1);
            let next_dim = self.shape.get(i + 1).copied().unwrap_or(1);
            if let Some(s) = strides.get_mut(i) {
                *s = next_stride * next_dim;
            }
        }

        let mut new_shape = self.shape.clone();
        if let Some(d) = new_shape.get_mut(dim) {
            *d = indices.len();
        }
        let inner_size = strides.get(dim).copied().unwrap_or(1);
        let outer_count = if dim == 0 { 1 } else {
            self.shape.iter().take(dim).product::<usize>()
        };

        // Upload indices to GPU
        let mut indices_dev = stream.alloc_zeros::<u32>(indices.len()).map_err(|e| {
            MLError::ModelError(format!("index_select indices alloc: {e}"))
        })?;
        stream.memcpy_htod(indices, &mut indices_dev).map_err(|e| {
            MLError::ModelError(format!("index_select indices HtoD: {e}"))
        })?;

        let kernels = super::elementwise::get_or_compile(stream)?;
        let out = kernels.gather(&self.data, &indices_dev, indices.len(), inner_size, outer_count, dim_size)?;
        Ok(Self { data: out, shape: new_shape })
    }

    /// Reshape without copying — returns error if total elements differ.
    pub fn reshape(self, new_shape: Vec<usize>) -> Result<Self, MLError> {
        let new_numel: usize = new_shape.iter().product();
        if new_numel != self.numel() {
            return Err(MLError::DimensionMismatch {
                expected: self.numel(),
                actual: new_numel,
            });
        }
        Ok(Self {
            data: self.data,
            shape: new_shape,
        })
    }
}

impl std::fmt::Debug for GpuTensor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "GpuTensor {{ shape: {:?}, numel: {} }}",
            self.shape,
            self.numel()
        )
    }
}

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

    #[test]
    fn test_zeros_shape() {
        let stream = cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream");
        let t = GpuTensor::zeros(&[4, 8], &stream).unwrap();
        assert_eq!(t.shape(), &[4, 8]);
        assert_eq!(t.numel(), 32);
    }

    #[test]
    fn test_roundtrip() {
        let stream = cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream");
        let host_data: Vec<f32> = (0..12).map(|i| i as f32).collect();
        let t = GpuTensor::from_host(&host_data, vec![3, 4], &stream).unwrap();
        let back = t.to_host(&stream).unwrap();
        assert_eq!(host_data, back);
    }

    #[test]
    fn test_shape_mismatch() {
        let stream = cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream");
        let host_data = vec![1.0_f32; 10];
        let result = GpuTensor::from_host(&host_data, vec![3, 4], &stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_reshape() {
        let stream = cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream");
        let t = GpuTensor::zeros(&[6, 4], &stream).unwrap();
        let t2 = t.reshape(vec![2, 12]).unwrap();
        assert_eq!(t2.shape(), &[2, 12]);
    }
}
