# Copyright (c) 2026 Andrew H. Bond
# Department of Computer Engineering, San Jose State University
# Licensed under the AGI-HPC Responsible AI License v1.0.

"""
CUDA Acceleration Backend using CuPy.

DEME V3 Sprint 12: GPU acceleration for tensor operations using CuPy,
providing 10x+ speedup for large moral tensors on NVIDIA GPUs.

Features:
- Drop-in NumPy replacement via CuPy
- Automatic memory management with memory pools
- Pinned memory for faster CPU-GPU transfers
- CUDA stream support for async operations
- Graceful fallback to CPU when CUDA unavailable

Requirements:
- NVIDIA GPU with CUDA support
- CuPy 12.0+ (pip install cupy-cuda12x)

Version: 3.0.0 (DEME V3)
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from .backend import (
    AccelerationBackend,
    DeviceInfo,
    DeviceType,
    TensorHandle,
)

logger = logging.getLogger(__name__)

# Try to import CuPy
try:
    import cupy as cp
    from cupy import cuda

    HAS_CUPY = True
except ImportError:
    HAS_CUPY = False
    cp = None  # type: ignore
    cuda = None  # type: ignore

# Numerical stability constant
EPSILON = 1e-10


class CUDABackend(AccelerationBackend):
    """
    CUDA-based acceleration backend using CuPy.

    Provides GPU-accelerated tensor operations for NVIDIA GPUs. CuPy
    provides a NumPy-compatible API that runs on CUDA, making it easy
    to port existing NumPy code to GPU.

    Key features:
    - Memory pool for efficient GPU memory management
    - Pinned memory for faster host-device transfers
    - CUDA streams for async operations
    - Multi-GPU support

    Example:
        backend = CUDABackend(device_id=0)
        if backend.is_available():
            tensor = backend.from_numpy(my_array)
            result = backend.multiply(tensor, 2.0)
            output = backend.to_numpy(result)
    """

    def __init__(
        self,
        device_id: int = 0,
        use_memory_pool: bool = True,
        use_pinned_memory: bool = True,
    ):
        """
        Initialize the CUDA backend.

        Args:
            device_id: CUDA device ID to use (default 0).
            use_memory_pool: Enable CuPy memory pool for efficiency.
            use_pinned_memory: Use pinned memory for faster transfers.
        """
        self._device_id = device_id
        self._use_memory_pool = use_memory_pool
        self._use_pinned_memory = use_pinned_memory
        self._device_info_cache: Dict[int, DeviceInfo] = {}
        self._initialized = False

        if HAS_CUPY:
            self._initialize()

    def _initialize(self) -> None:
        """Initialize CUDA device and memory pools."""
        if not HAS_CUPY:
            return

        try:
            # Set the device
            with cuda.Device(self._device_id):
                # Configure memory pool
                if self._use_memory_pool:
                    # Access memory pools to ensure they are initialized
                    # CuPy uses these pools by default for efficient allocation
                    cp.get_default_memory_pool()
                    cp.get_default_pinned_memory_pool()

                self._initialized = True
                logger.debug(f"CUDA backend initialized on device {self._device_id}")

        except Exception as e:
            logger.warning(f"Failed to initialize CUDA backend: {e}")
            self._initialized = False

    @property
    def name(self) -> str:
        return "cuda"

    @property
    def device_type(self) -> DeviceType:
        return DeviceType.CUDA

    def is_available(self) -> bool:
        """Check if CUDA is available."""
        if not HAS_CUPY:
            return False

        try:
            device_count = cuda.runtime.getDeviceCount()
            return device_count > 0 and self._device_id < device_count
        except Exception:
            return False

    def get_device_info(self, device_id: int = 0) -> DeviceInfo:
        """Get information about a CUDA device."""
        if device_id in self._device_info_cache:
            return self._device_info_cache[device_id]

        if not self.is_available():
            return DeviceInfo(
                device_type=DeviceType.CUDA,
                device_id=device_id,
                name="CUDA (unavailable)",
                is_available=False,
            )

        try:
            with cuda.Device(device_id):
                props = cuda.runtime.getDeviceProperties(device_id)
                mem_info = cuda.runtime.memGetInfo()

                info = DeviceInfo(
                    device_type=DeviceType.CUDA,
                    device_id=device_id,
                    name=(
                        props["name"].decode()
                        if isinstance(props["name"], bytes)
                        else str(props["name"])
                    ),
                    memory_total=props["totalGlobalMem"],
                    memory_available=mem_info[0],  # Free memory
                    compute_capability=(props["major"], props["minor"]),
                    is_available=True,
                    properties={
                        "multiProcessorCount": props["multiProcessorCount"],
                        "maxThreadsPerBlock": props["maxThreadsPerBlock"],
                        "maxThreadsPerMultiProcessor": props[
                            "maxThreadsPerMultiProcessor"
                        ],
                        "warpSize": props["warpSize"],
                        "clockRate": props["clockRate"],
                        "memoryClockRate": props["memoryClockRate"],
                        "l2CacheSize": props["l2CacheSize"],
                        "cupy_version": cp.__version__,
                    },
                )
                self._device_info_cache[device_id] = info
                return info

        except Exception as e:
            logger.warning(f"Failed to get device info for device {device_id}: {e}")
            return DeviceInfo(
                device_type=DeviceType.CUDA,
                device_id=device_id,
                name="CUDA (error)",
                is_available=False,
                properties={"error": str(e)},
            )

    def get_all_devices(self) -> List[DeviceInfo]:
        """Get information about all CUDA devices."""
        if not HAS_CUPY:
            return []

        try:
            device_count = cuda.runtime.getDeviceCount()
            return [self.get_device_info(i) for i in range(device_count)]
        except Exception:
            return []

    def _ensure_device(self) -> None:
        """Ensure we're on the correct device."""
        if HAS_CUPY:
            cuda.Device(self._device_id).use()

    def _to_cupy(self, array: np.ndarray) -> Any:
        """Convert NumPy array to CuPy array."""
        self._ensure_device()

        if self._use_pinned_memory:
            # Use pinned memory for faster transfer
            try:
                pinned = cp.cuda.alloc_pinned_memory(array.nbytes)
                pinned_array = np.frombuffer(
                    pinned, dtype=array.dtype, count=array.size
                ).reshape(array.shape)
                pinned_array[...] = array
                return cp.asarray(pinned_array)
            except Exception:
                # Fall back to regular transfer
                return cp.asarray(array)
        else:
            return cp.asarray(array)

    def _from_cupy(self, cupy_array: Any) -> np.ndarray:
        """Convert CuPy array to NumPy array."""
        return cp.asnumpy(cupy_array)

    # -------------------------------------------------------------------------
    # Tensor Creation
    # -------------------------------------------------------------------------

    def from_numpy(
        self,
        array: np.ndarray,
        device_id: int = 0,
    ) -> TensorHandle:
        """Create a GPU tensor from a NumPy array."""
        if not self.is_available():
            raise RuntimeError("CUDA backend is not available")

        self._ensure_device()

        # Convert to float64 and transfer to GPU
        arr_np = np.asarray(array, dtype=np.float64)
        arr_gpu = self._to_cupy(arr_np)

        return TensorHandle(
            backend_name=self.name,
            device_type=self.device_type,
            device_id=device_id,
            shape=arr_gpu.shape,
            dtype="float64",
            _native_tensor=arr_gpu,
            is_sparse=False,
        )

    def to_numpy(self, handle: TensorHandle) -> np.ndarray:
        """Copy GPU tensor back to NumPy array."""
        if handle._native_tensor is None:
            raise ValueError("TensorHandle has no underlying tensor")

        return self._from_cupy(handle._native_tensor)

    def zeros(
        self,
        shape: Tuple[int, ...],
        dtype: str = "float64",
        device_id: int = 0,
    ) -> TensorHandle:
        """Create a zero-filled GPU tensor."""
        self._ensure_device()
        arr = cp.zeros(shape, dtype=cp.dtype(dtype))
        return TensorHandle(
            backend_name=self.name,
            device_type=self.device_type,
            device_id=device_id,
            shape=shape,
            dtype=dtype,
            _native_tensor=arr,
            is_sparse=False,
        )

    def ones(
        self,
        shape: Tuple[int, ...],
        dtype: str = "float64",
        device_id: int = 0,
    ) -> TensorHandle:
        """Create a one-filled GPU tensor."""
        self._ensure_device()
        arr = cp.ones(shape, dtype=cp.dtype(dtype))
        return TensorHandle(
            backend_name=self.name,
            device_type=self.device_type,
            device_id=device_id,
            shape=shape,
            dtype=dtype,
            _native_tensor=arr,
            is_sparse=False,
        )

    def full(
        self,
        shape: Tuple[int, ...],
        fill_value: float,
        dtype: str = "float64",
        device_id: int = 0,
    ) -> TensorHandle:
        """Create a constant-filled GPU tensor."""
        self._ensure_device()
        arr = cp.full(shape, fill_value, dtype=cp.dtype(dtype))
        return TensorHandle(
            backend_name=self.name,
            device_type=self.device_type,
            device_id=device_id,
            shape=shape,
            dtype=dtype,
            _native_tensor=arr,
            is_sparse=False,
        )

    # -------------------------------------------------------------------------
    # Helper Methods
    # -------------------------------------------------------------------------

    def _get_tensor(self, handle: TensorHandle) -> Any:
        """Extract CuPy array from handle."""
        if handle._native_tensor is None:
            raise ValueError("TensorHandle has no underlying tensor")
        return handle._native_tensor

    def _wrap_result(
        self,
        result: Any,
        device_id: int = 0,
    ) -> TensorHandle:
        """Wrap CuPy array result in a TensorHandle."""
        return TensorHandle(
            backend_name=self.name,
            device_type=self.device_type,
            device_id=device_id,
            shape=result.shape,
            dtype=str(result.dtype),
            _native_tensor=result,
            is_sparse=False,
        )

    # -------------------------------------------------------------------------
    # Element-wise Operations
    # -------------------------------------------------------------------------

    def add(
        self,
        a: TensorHandle,
        b: Union[TensorHandle, float],
    ) -> TensorHandle:
        """Element-wise addition on GPU."""
        a_arr = self._get_tensor(a)
        if isinstance(b, TensorHandle):
            b_arr = self._get_tensor(b)
            result = cp.add(a_arr, b_arr)
        else:
            result = cp.add(a_arr, b)
        return self._wrap_result(result, a.device_id)

    def subtract(
        self,
        a: TensorHandle,
        b: Union[TensorHandle, float],
    ) -> TensorHandle:
        """Element-wise subtraction on GPU."""
        a_arr = self._get_tensor(a)
        if isinstance(b, TensorHandle):
            b_arr = self._get_tensor(b)
            result = cp.subtract(a_arr, b_arr)
        else:
            result = cp.subtract(a_arr, b)
        return self._wrap_result(result, a.device_id)

    def multiply(
        self,
        a: TensorHandle,
        b: Union[TensorHandle, float],
    ) -> TensorHandle:
        """Element-wise multiplication on GPU."""
        a_arr = self._get_tensor(a)
        if isinstance(b, TensorHandle):
            b_arr = self._get_tensor(b)
            result = cp.multiply(a_arr, b_arr)
        else:
            result = cp.multiply(a_arr, b)
        return self._wrap_result(result, a.device_id)

    def divide(
        self,
        a: TensorHandle,
        b: Union[TensorHandle, float],
        safe: bool = True,
    ) -> TensorHandle:
        """Element-wise division on GPU with safe handling."""
        a_arr = self._get_tensor(a)
        if isinstance(b, TensorHandle):
            b_arr = self._get_tensor(b)
            if safe:
                result = cp.where(
                    cp.abs(b_arr) < EPSILON,
                    1.0,
                    a_arr / b_arr,
                )
            else:
                result = cp.divide(a_arr, b_arr)
        else:
            if safe and abs(b) < EPSILON:
                result = cp.ones_like(a_arr)
            else:
                result = cp.divide(a_arr, b)
        return self._wrap_result(result, a.device_id)

    def clip(
        self,
        a: TensorHandle,
        min_val: float,
        max_val: float,
    ) -> TensorHandle:
        """Clip tensor values on GPU."""
        result = cp.clip(self._get_tensor(a), min_val, max_val)
        return self._wrap_result(result, a.device_id)

    def abs(self, a: TensorHandle) -> TensorHandle:
        """Element-wise absolute value on GPU."""
        result = cp.abs(self._get_tensor(a))
        return self._wrap_result(result, a.device_id)

    def sqrt(self, a: TensorHandle) -> TensorHandle:
        """Element-wise square root on GPU."""
        result = cp.sqrt(self._get_tensor(a))
        return self._wrap_result(result, a.device_id)

    def exp(self, a: TensorHandle) -> TensorHandle:
        """Element-wise exponential on GPU."""
        result = cp.exp(self._get_tensor(a))
        return self._wrap_result(result, a.device_id)

    def log(self, a: TensorHandle, safe: bool = True) -> TensorHandle:
        """Element-wise natural logarithm on GPU."""
        a_arr = self._get_tensor(a)
        if safe:
            a_arr = cp.clip(a_arr, EPSILON, None)
        result = cp.log(a_arr)
        return self._wrap_result(result, a.device_id)

    # -------------------------------------------------------------------------
    # Reduction Operations
    # -------------------------------------------------------------------------

    def sum(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
        keepdims: bool = False,
    ) -> TensorHandle:
        """Sum reduction on GPU."""
        result = cp.sum(self._get_tensor(a), axis=axis, keepdims=keepdims)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result)
        return self._wrap_result(result, a.device_id)

    def mean(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
        keepdims: bool = False,
    ) -> TensorHandle:
        """Mean reduction on GPU."""
        result = cp.mean(self._get_tensor(a), axis=axis, keepdims=keepdims)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result)
        return self._wrap_result(result, a.device_id)

    def min(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
        keepdims: bool = False,
    ) -> TensorHandle:
        """Minimum reduction on GPU."""
        result = cp.min(self._get_tensor(a), axis=axis, keepdims=keepdims)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result)
        return self._wrap_result(result, a.device_id)

    def max(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
        keepdims: bool = False,
    ) -> TensorHandle:
        """Maximum reduction on GPU."""
        result = cp.max(self._get_tensor(a), axis=axis, keepdims=keepdims)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result)
        return self._wrap_result(result, a.device_id)

    def argmin(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
    ) -> TensorHandle:
        """Index of minimum on GPU."""
        result = cp.argmin(self._get_tensor(a), axis=axis)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result, dtype=cp.int64)
        return self._wrap_result(result, a.device_id)

    def argmax(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
    ) -> TensorHandle:
        """Index of maximum on GPU."""
        result = cp.argmax(self._get_tensor(a), axis=axis)
        if not isinstance(result, cp.ndarray):
            result = cp.array(result, dtype=cp.int64)
        return self._wrap_result(result, a.device_id)

    # -------------------------------------------------------------------------
    # Linear Algebra Operations
    # -------------------------------------------------------------------------

    def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
        """Dot product on GPU."""
        result = cp.dot(self._get_tensor(a), self._get_tensor(b))
        if not isinstance(result, cp.ndarray):
            result = cp.array([result])
        return self._wrap_result(result, a.device_id)

    def tensordot(
        self,
        a: TensorHandle,
        b: TensorHandle,
        axes: Union[int, Tuple[List[int], List[int]]],
    ) -> TensorHandle:
        """Tensor contraction on GPU."""
        result = cp.tensordot(self._get_tensor(a), self._get_tensor(b), axes=axes)
        return self._wrap_result(result, a.device_id)

    def norm(
        self,
        a: TensorHandle,
        ord: Optional[Union[int, float, str]] = None,
        axis: Optional[int] = None,
    ) -> TensorHandle:
        """Tensor norm on GPU."""
        result = cp.linalg.norm(self._get_tensor(a), ord=ord, axis=axis)
        if not isinstance(result, cp.ndarray):
            result = cp.array([result])
        return self._wrap_result(result, a.device_id)

    # -------------------------------------------------------------------------
    # Shape Operations
    # -------------------------------------------------------------------------

    def reshape(
        self,
        a: TensorHandle,
        shape: Tuple[int, ...],
    ) -> TensorHandle:
        """Reshape tensor on GPU."""
        result = cp.reshape(self._get_tensor(a), shape)
        return self._wrap_result(result, a.device_id)

    def transpose(
        self,
        a: TensorHandle,
        axes: Optional[Tuple[int, ...]] = None,
    ) -> TensorHandle:
        """Transpose tensor on GPU."""
        result = cp.transpose(self._get_tensor(a), axes=axes)
        return self._wrap_result(result, a.device_id)

    def expand_dims(
        self,
        a: TensorHandle,
        axis: int,
    ) -> TensorHandle:
        """Insert new axis on GPU."""
        result = cp.expand_dims(self._get_tensor(a), axis=axis)
        return self._wrap_result(result, a.device_id)

    def squeeze(
        self,
        a: TensorHandle,
        axis: Optional[int] = None,
    ) -> TensorHandle:
        """Remove singleton dimensions on GPU."""
        result = cp.squeeze(self._get_tensor(a), axis=axis)
        return self._wrap_result(result, a.device_id)

    def broadcast_to(
        self,
        a: TensorHandle,
        shape: Tuple[int, ...],
    ) -> TensorHandle:
        """Broadcast tensor on GPU."""
        result = cp.broadcast_to(self._get_tensor(a), shape).copy()
        return self._wrap_result(result, a.device_id)

    # -------------------------------------------------------------------------
    # Concatenation and Stacking
    # -------------------------------------------------------------------------

    def concatenate(
        self,
        tensors: List[TensorHandle],
        axis: int = 0,
    ) -> TensorHandle:
        """Concatenate tensors on GPU."""
        arrays = [self._get_tensor(t) for t in tensors]
        result = cp.concatenate(arrays, axis=axis)
        device_id = tensors[0].device_id if tensors else 0
        return self._wrap_result(result, device_id)

    def stack(
        self,
        tensors: List[TensorHandle],
        axis: int = 0,
    ) -> TensorHandle:
        """Stack tensors on GPU."""
        arrays = [self._get_tensor(t) for t in tensors]
        result = cp.stack(arrays, axis=axis)
        device_id = tensors[0].device_id if tensors else 0
        return self._wrap_result(result, device_id)

    # -------------------------------------------------------------------------
    # Sparse Operations
    # -------------------------------------------------------------------------

    def to_sparse_coo(
        self,
        a: TensorHandle,
        fill_value: float = 0.0,
    ) -> Tuple[TensorHandle, TensorHandle, Tuple[int, ...]]:
        """Convert dense to sparse COO on GPU."""
        data = self._get_tensor(a)
        shape = data.shape

        if fill_value == 0.0:
            mask = cp.abs(data) > EPSILON
        else:
            mask = cp.abs(data - fill_value) > EPSILON

        coords = cp.argwhere(mask).astype(cp.int32)
        values = data[mask].astype(cp.float64)

        return (
            self._wrap_result(coords, a.device_id),
            self._wrap_result(values, a.device_id),
            shape,
        )

    def from_sparse_coo(
        self,
        coords: TensorHandle,
        values: TensorHandle,
        shape: Tuple[int, ...],
        fill_value: float = 0.0,
    ) -> TensorHandle:
        """Convert sparse COO to dense on GPU."""
        coords_gpu = self._get_tensor(coords)
        values_gpu = self._get_tensor(values)

        # Ensure coordinates are integers for indexing
        coords_int = coords_gpu.astype(cp.int64)

        dense = cp.full(shape, fill_value, dtype=cp.float64)
        if len(coords_int) > 0:
            idx = tuple(coords_int[:, i] for i in range(len(shape)))
            dense[idx] = values_gpu

        return self._wrap_result(dense, coords.device_id)

    # -------------------------------------------------------------------------
    # Advanced Operations
    # -------------------------------------------------------------------------

    def sort(
        self,
        a: TensorHandle,
        axis: int = -1,
    ) -> TensorHandle:
        """Sort tensor on GPU."""
        result = cp.sort(self._get_tensor(a), axis=axis)
        return self._wrap_result(result, a.device_id)

    def argsort(
        self,
        a: TensorHandle,
        axis: int = -1,
    ) -> TensorHandle:
        """Argsort on GPU."""
        result = cp.argsort(self._get_tensor(a), axis=axis)
        return self._wrap_result(result.astype(cp.int64), a.device_id)

    def where(
        self,
        condition: TensorHandle,
        x: TensorHandle,
        y: TensorHandle,
    ) -> TensorHandle:
        """Conditional selection on GPU."""
        cond_gpu = self._get_tensor(condition)
        x_gpu = self._get_tensor(x)
        y_gpu = self._get_tensor(y)
        result = cp.where(cond_gpu, x_gpu, y_gpu)
        return self._wrap_result(result, x.device_id)

    def einsum(
        self,
        subscripts: str,
        *operands: TensorHandle,
    ) -> TensorHandle:
        """Einstein summation on GPU."""
        arrays = [self._get_tensor(op) for op in operands]
        result = cp.einsum(subscripts, *arrays, optimize=True)
        device_id = operands[0].device_id if operands else 0
        return self._wrap_result(result, device_id)

    # -------------------------------------------------------------------------
    # Memory Management
    # -------------------------------------------------------------------------

    def synchronize(self, device_id: int = 0) -> None:
        """Synchronize CUDA device."""
        if HAS_CUPY:
            with cuda.Device(device_id):
                cuda.Stream.null.synchronize()

    def get_memory_info(self, device_id: int = 0) -> Tuple[int, int]:
        """Get GPU memory info (used, total)."""
        if not HAS_CUPY:
            return (0, 0)

        try:
            with cuda.Device(device_id):
                free, total = cuda.runtime.memGetInfo()
                used = total - free
                return (used, total)
        except Exception:
            return (0, 0)

    def clear_cache(self, device_id: int = 0) -> None:
        """Clear GPU memory cache."""
        if HAS_CUPY:
            mempool = cp.get_default_memory_pool()
            mempool.free_all_blocks()
            pinned_mempool = cp.get_default_pinned_memory_pool()
            pinned_mempool.free_all_blocks()

    # -------------------------------------------------------------------------
    # Ethics-Specific Optimizations
    # -------------------------------------------------------------------------

    def moral_contraction(
        self,
        tensor: TensorHandle,
        weights: TensorHandle,
        axis: int,
        normalize: bool = True,
    ) -> TensorHandle:
        """
        Optimized moral tensor contraction on GPU.

        Uses CuPy's tensordot for efficient weighted reduction.
        """
        t_arr = self._get_tensor(tensor)
        w_arr = self._get_tensor(weights)

        if normalize:
            w_sum = cp.sum(w_arr)
            if w_sum > EPSILON:
                w_arr = w_arr / w_sum
            else:
                w_arr = cp.ones_like(w_arr) / len(w_arr)

        result = cp.tensordot(t_arr, w_arr, axes=([axis], [0]))
        result = cp.clip(result, 0.0, 1.0)

        return self._wrap_result(result, tensor.device_id)

    def batch_gini(self, values: TensorHandle, axis: int = -1) -> TensorHandle:
        """
        Batched Gini coefficient computation on GPU.

        Optimized for parallel computation across tensor slices.
        """
        v = self._get_tensor(values)

        # Move target axis to last position
        v = cp.moveaxis(v, axis, -1)
        original_shape = v.shape[:-1]
        n = v.shape[-1]

        if n <= 1:
            return self._wrap_result(
                cp.zeros(original_shape, dtype=cp.float64), values.device_id
            )

        # Flatten batch dimensions
        batch_size = int(cp.prod(cp.array(original_shape))) if original_shape else 1
        v_flat = v.reshape(batch_size, n)

        # Sort each row (parallel on GPU)
        v_sorted = cp.sort(v_flat, axis=1)

        # Compute Gini for each row
        indices = cp.arange(1, n + 1, dtype=cp.float64)
        gini = (
            2 * cp.sum(indices * v_sorted, axis=1) - (n + 1) * cp.sum(v_sorted, axis=1)
        ) / (n * cp.sum(v_sorted, axis=1) + EPSILON)

        # Handle edge cases
        gini = cp.where(cp.sum(v_sorted, axis=1) < EPSILON, 0.0, gini)

        if original_shape:
            gini = gini.reshape(original_shape)

        return self._wrap_result(gini, values.device_id)

    # -------------------------------------------------------------------------
    # Batched Operations
    # -------------------------------------------------------------------------

    def batch_matmul(
        self,
        a: TensorHandle,
        b: TensorHandle,
    ) -> TensorHandle:
        """
        Batched matrix multiplication on GPU.

        Efficiently handles batched matrix operations common in
        multi-agent ethics computations.

        Args:
            a: Tensor of shape (..., M, K)
            b: Tensor of shape (..., K, N)

        Returns:
            Tensor of shape (..., M, N)
        """
        a_arr = self._get_tensor(a)
        b_arr = self._get_tensor(b)
        result = cp.matmul(a_arr, b_arr)
        return self._wrap_result(result, a.device_id)

    def batch_trace(self, a: TensorHandle) -> TensorHandle:
        """
        Batched trace computation on GPU.

        Args:
            a: Tensor of shape (..., N, N)

        Returns:
            Tensor of shape (...) with traces
        """
        a_arr = self._get_tensor(a)
        result = cp.trace(a_arr, axis1=-2, axis2=-1)
        return self._wrap_result(result, a.device_id)


# Check for CUDA availability at module load
def cuda_is_available() -> bool:
    """Check if CUDA is available."""
    if not HAS_CUPY:
        return False
    try:
        return cuda.runtime.getDeviceCount() > 0
    except Exception:
        return False


def get_cuda_backend(device_id: int = 0) -> CUDABackend:
    """Get a CUDA backend instance for the specified device."""
    return CUDABackend(device_id=device_id)


__all__ = [
    "CUDABackend",
    "cuda_is_available",
    "get_cuda_backend",
    "HAS_CUPY",
]
