"""
Unified array operations with automatic backend dispatch.

Functions in this module automatically detect the backend from input arrays
and dispatch to the appropriate implementation. Simple operations use inline
if/else, while complex operations use a dispatch table to numpy_ops/torch_ops.
"""

from __future__ import annotations

import types

import numpy as np

from .core import get_backend, is_torch, Backend, Array


# =============================================================================
# Backend Validation
# =============================================================================

def _check_backends(*arrays: Array) -> None:
    """
    Check that all arrays have the same backend.

    Raises:
        TypeError: If backends don't match.
    """
    if len(arrays) < 2:
        return

    first_is_torch = is_torch(arrays[0])
    for i, arr in enumerate(arrays[1:], 1):
        if is_torch(arr) != first_is_torch:
            first_backend = "torch" if first_is_torch else "numpy"
            other_backend = "torch" if is_torch(arr) else "numpy"
            raise TypeError(
                f"Backend mismatch: arrays[0] is {first_backend}, "
                f"arrays[{i}] is {other_backend}. Convert to same backend first."
            )


# =============================================================================
# Dispatch Table
# =============================================================================

def _get_ops(arr: Array) -> "types.ModuleType":
    """Get the appropriate ops module for the array's backend."""
    if get_backend(arr) == Backend.TORCH:
        from . import torch_ops
        return torch_ops
    from . import numpy_ops
    return numpy_ops


# =============================================================================
# Scatter Operations
# =============================================================================

def scatter_sum(src: Array, index: Array, dim_size: int) -> Array:
    """
    Sum values into an output array at specified indices.

    Args:
        src: Source array of shape (N, ...).
        index: Index array of shape (N,) with values in [0, dim_size).
        dim_size: Size of output first dimension.

    Returns:
        Array of shape (dim_size, ...) with summed values.
    """
    _check_backends(src, index)
    return _get_ops(src).scatter_sum(src, index, dim_size)


def scatter_mean(src: Array, index: Array, dim_size: int) -> Array:
    """
    Average values into an output array at specified indices.

    Args:
        src: Source array of shape (N, ...).
        index: Index array of shape (N,) with values in [0, dim_size).
        dim_size: Size of output first dimension.

    Returns:
        Array of shape (dim_size, ...) with averaged values.
    """
    _check_backends(src, index)
    return _get_ops(src).scatter_mean(src, index, dim_size)


def scatter_max(src: Array, index: Array, dim_size: int) -> tuple[Array, Array | None]:
    """
    Maximum values at specified indices.

    Args:
        src: Source array of shape (N, ...).
        index: Index array of shape (N,) with values in [0, dim_size).
        dim_size: Size of output first dimension.

    Returns:
        Tuple of (max_values, argmax_indices). argmax_indices may be None.
    """
    _check_backends(src, index)
    return _get_ops(src).scatter_max(src, index, dim_size)


def scatter_min(src: Array, index: Array, dim_size: int) -> tuple[Array, Array | None]:
    """
    Minimum values at specified indices.

    Args:
        src: Source array of shape (N, ...).
        index: Index array of shape (N,) with values in [0, dim_size).
        dim_size: Size of output first dimension.

    Returns:
        Tuple of (min_values, argmin_indices). argmin_indices may be None.
    """
    _check_backends(src, index)
    return _get_ops(src).scatter_min(src, index, dim_size)


# =============================================================================
# Array Operations
# =============================================================================

def repeat_interleave(arr: Array, repeats: Array) -> Array:
    """
    Repeat elements of an array along the first axis.

    Args:
        arr: Array to repeat elements from.
        repeats: Number of times to repeat each element.

    Returns:
        Array with repeated elements.
    """
    _check_backends(arr, repeats)
    return _get_ops(arr).repeat_interleave(arr, repeats)


def cdist(x1: Array, x2: Array) -> Array:
    """
    Compute pairwise Euclidean distances.

    Args:
        x1: Array of shape (M, D).
        x2: Array of shape (N, D).

    Returns:
        Distance matrix of shape (M, N).
    """
    _check_backends(x1, x2)
    return _get_ops(x1).cdist(x1, x2)


def cat(arrays: list, axis: int = 0) -> Array:
    """
    Concatenate arrays along an axis.

    Args:
        arrays: List of arrays to concatenate.
        axis: Axis along which to concatenate.

    Returns:
        Concatenated array.
    """
    if len(arrays) == 0:
        raise ValueError("Cannot concatenate empty list")

    _check_backends(*arrays)
    if is_torch(arrays[0]):
        import torch
        return torch.cat(arrays, dim=axis)
    return np.concatenate(arrays, axis=axis)


def stack(arrays: list, axis: int = 0) -> Array:
    """
    Stack arrays along a new axis.

    Args:
        arrays: List of arrays to stack.
        axis: Axis along which to stack.

    Returns:
        Stacked array with new dimension.
    """
    if len(arrays) == 0:
        raise ValueError("Cannot stack empty list")

    _check_backends(*arrays)
    if is_torch(arrays[0]):
        import torch
        return torch.stack(arrays, dim=axis)
    return np.stack(arrays, axis=axis)


def nonzero(condition: Array) -> tuple[Array, ...]:
    """
    Return indices where condition is True (nonzero).

    Args:
        condition: Boolean array.

    Returns:
        Tuple of arrays, one per dimension, containing indices where
        condition is True. For 1D input, returns a single-element tuple.

    Note:
        This is equivalent to numpy.where(condition) with a single argument,
        or torch.where(condition). For the ternary version that selects
        between two arrays, use where(condition, x, y).
    """
    if is_torch(condition):
        import torch
        return torch.where(condition)
    return np.where(condition)


def argwhere(condition: Array) -> Array:
    """
    Return indices where condition is True as (N, ndim) array.

    Args:
        condition: Boolean array.

    Returns:
        (N, ndim) array where each row contains the indices of a True element.

    Note:
        Unlike nonzero() which returns a tuple of arrays, this returns a
        single stacked array, similar to numpy.argwhere or torch.nonzero.
    """
    if is_torch(condition):
        import torch
        return torch.nonzero(condition)
    return np.argwhere(condition)


def triu(arr: Array, diagonal: int = 0) -> Array:
    """
    Return upper triangular part of an array.

    Args:
        arr: Input 2D array.
        diagonal: Diagonal offset. 0 = main diagonal, 1 = above main, etc.

    Returns:
        Array with elements below the diagonal zeroed.
    """
    if is_torch(arr):
        import torch
        return torch.triu(arr, diagonal=diagonal)
    return np.triu(arr, k=diagonal)


def argmax(arr: Array, axis: int | None = None) -> Array:
    """
    Return indices of maximum values along an axis.

    Args:
        arr: Input array.
        axis: Axis along which to find argmax. If None, returns index
              into flattened array.

    Returns:
        Array of indices.
    """
    if is_torch(arr):
        import torch
        return torch.argmax(arr, dim=axis)
    return np.argmax(arr, axis=axis)


def any(arr: Array, axis: int | None = None) -> Array:
    """
    Test whether any element along an axis is True.

    Args:
        arr: Input array.
        axis: Axis along which to test. If None, tests all elements.

    Returns:
        Boolean array (or scalar if axis is None).
    """
    if is_torch(arr):
        import torch
        if axis is None:
            return torch.any(arr)
        return torch.any(arr, dim=axis)
    return np.any(arr, axis=axis)


def multiply(a: Array, b: Array) -> Array:
    """
    Element-wise multiplication.

    Args:
        a: First array.
        b: Second array.

    Returns:
        Element-wise product.
    """
    _check_backends(a, b)
    return _get_ops(a).multiply(a, b)


def sign(arr: Array) -> Array:
    """
    Element-wise sign (-1, 0, or 1).

    Args:
        arr: Input array.

    Returns:
        Array with sign of each element.
    """
    return _get_ops(arr).sign(arr)


def arange(n: int, like: Array) -> Array:
    """
    Create an integer range [0, n) with same backend/device as reference array.

    Args:
        n: Length of the range.
        like: Reference array to match backend and device.

    Returns:
        Integer array [0, 1, ..., n-1] on same backend/device as `like`.
    """
    ops = _get_ops(like)
    if get_backend(like) == Backend.TORCH:
        return ops.arange(n, device=like.device)
    return ops.arange(n)


def cumsum(arr: Array, dim: int = 0) -> Array:
    """
    Compute cumulative sum along a dimension.

    Args:
        arr: Input array.
        dim: Dimension along which to compute cumsum.

    Returns:
        Array with cumulative sums.
    """
    if is_torch(arr):
        import torch
        return torch.cumsum(arr, dim=dim)
    return np.cumsum(arr, axis=dim)


# =============================================================================
# Backend Conversion
# =============================================================================

def to_backend(arr: np.ndarray, like: Array) -> Array:
    """
    Convert a numpy array to match the backend of 'like'.

    Args:
        arr: NumPy array to convert.
        like: Template array whose backend to match.

    Returns:
        Array in the same backend as 'like'. If 'like' is torch,
        returns a tensor on the same device.
    """
    if is_torch(like):
        import torch
        result = torch.from_numpy(np.ascontiguousarray(arr))
        if hasattr(like, 'device'):
            result = result.to(like.device)
        return result
    return arr


def convert_backend(arr: Array, like: Array) -> Array:
    """
    Convert arr to match the backend and device of 'like'.

    More general than to_backend - works with both numpy and torch inputs.
    When both are torch tensors, also aligns device.

    Args:
        arr: Array to convert.
        like: Template array for backend and device detection.

    Returns:
        Array in the same backend and device as 'like'.
    """
    if is_torch(like):
        if not is_torch(arr):
            import torch
            return torch.from_numpy(np.asarray(arr)).to(like.device)
        # Both are torch - also align device
        if arr.device != like.device:
            return arr.to(like.device)
        return arr
    else:
        if is_torch(arr):
            return arr.detach().cpu().numpy()
        return np.asarray(arr)


# =============================================================================
# Linear Algebra
# =============================================================================

def eigh(arr: Array) -> tuple[Array, Array]:
    """
    Eigenvalue decomposition of a symmetric/Hermitian matrix.

    Args:
        arr: Symmetric matrix.

    Returns:
        Tuple of (eigenvalues, eigenvectors) in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.linalg.eigh(arr)
    return np.linalg.eigh(arr)


def det(arr: Array) -> Array:
    """
    Matrix determinant.

    Args:
        arr: Square matrix.

    Returns:
        Determinant value in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.linalg.det(arr)
    return np.linalg.det(arr)


def svd(arr: Array) -> tuple[Array, Array, Array]:
    """
    Full singular value decomposition.

    Args:
        arr: Input matrix.

    Returns:
        Tuple of (U, S, Vh) matrices in original backend.
    """
    if is_torch(arr):
        import torch
        # MPS doesn't support SVD - fallback to CPU
        if arr.device.type == "mps":
            U, S, Vh = torch.linalg.svd(arr.cpu())
            return U.to(arr.device), S.to(arr.device), Vh.to(arr.device)
        return torch.linalg.svd(arr)
    return np.linalg.svd(arr)


def svdvals(arr: Array) -> Array:
    """
    Singular values of a matrix.

    Args:
        arr: Input matrix.

    Returns:
        Singular values in original backend.
    """
    if is_torch(arr):
        import torch
        # MPS doesn't support SVD - fallback to CPU
        if arr.device.type == "mps":
            return torch.linalg.svdvals(arr.cpu()).to(arr.device)
        return torch.linalg.svdvals(arr)
    return np.linalg.svd(arr, compute_uv=False)


def pinv(arr: Array, rtol: float = 1e-15) -> Array:
    """
    Moore-Penrose pseudo-inverse of a matrix.

    Args:
        arr: Input matrix.
        rtol: Relative tolerance for singular values.

    Returns:
        Pseudo-inverse matrix in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.linalg.pinv(arr, rtol=rtol)
    return np.linalg.pinv(arr, rcond=rtol)


def diag(arr: Array) -> Array:
    """
    Create a diagonal matrix from a 1D array, or extract diagonal from 2D.

    Args:
        arr: 1D array to place on diagonal, or 2D matrix to extract from.

    Returns:
        If 1D input: 2D diagonal matrix.
        If 2D input: 1D diagonal vector.
    """
    if is_torch(arr):
        import torch
        return torch.diag(arr)
    return np.diag(arr)


def diagonal(arr: Array) -> Array:
    """
    Extract the diagonal of a 2D matrix.

    Args:
        arr: 2D matrix.

    Returns:
        1D array containing the diagonal elements.
    """
    if is_torch(arr):
        import torch
        return torch.diagonal(arr)
    return np.diagonal(arr)


# =============================================================================
# Array Creation (backend-aware)
# =============================================================================

def zeros(size: int, *, like: Array, dtype: str = 'int64') -> Array:
    """
    Create a zeros array matching the backend of 'like'.

    Args:
        size: Length of array.
        like: Template array for backend detection.
        dtype: Data type ('int64', 'float32', 'bool').

    Returns:
        Zeros array in the same backend as 'like'.
    """
    if is_torch(like):
        import torch
        torch_dtype = {'int64': torch.long, 'float32': torch.float32, 'bool': torch.bool}[dtype]
        return torch.zeros(size, dtype=torch_dtype, device=getattr(like, 'device', None))
    np_dtype = {'int64': np.int64, 'float32': np.float32, 'bool': bool}[dtype]
    return np.zeros(size, dtype=np_dtype)


def ones(size: int, *, like: Array, dtype: str = 'int64') -> Array:
    """
    Create a ones array matching the backend of 'like'.

    Args:
        size: Length of array.
        like: Template array for backend detection.
        dtype: Data type ('int64', 'float32').

    Returns:
        Ones array in the same backend as 'like'.
    """
    if is_torch(like):
        import torch
        torch_dtype = {'int64': torch.long, 'float32': torch.float32}[dtype]
        return torch.ones(size, dtype=torch_dtype, device=getattr(like, 'device', None))
    np_dtype = {'int64': np.int64, 'float32': np.float32}[dtype]
    return np.ones(size, dtype=np_dtype)


def array(data: list, *, like: Array, dtype: str = 'int64') -> Array:
    """
    Create an array from data matching the backend of 'like'.

    Args:
        data: List of values.
        like: Template array for backend detection.
        dtype: Data type ('int64', 'float32').

    Returns:
        Array in the same backend as 'like'.
    """
    if is_torch(like):
        import torch
        torch_dtype = {'int64': torch.long, 'float32': torch.float32}[dtype]
        return torch.tensor(data, dtype=torch_dtype, device=getattr(like, 'device', None))
    np_dtype = {'int64': np.int64, 'float32': np.float32}[dtype]
    return np.array(data, dtype=np_dtype)


def empty(size: int | tuple, *, like: Array, dtype: str | None = None) -> Array:
    """
    Create an empty (uninitialized) array matching backend/device of 'like'.

    Args:
        size: Shape of the array (int for 1D, tuple for nD).
        like: Template array to match backend and device.
        dtype: Data type ('int64', 'float32', 'bool'). If None, uses like's dtype.

    Returns:
        Empty array in the same backend/device as 'like'.
    """
    if is_torch(like):
        import torch
        if dtype is None:
            torch_dtype = like.dtype
        else:
            torch_dtype = {'int64': torch.long, 'float32': torch.float32, 'bool': torch.bool}[dtype]
        return torch.empty(size, dtype=torch_dtype, device=getattr(like, 'device', None))
    if dtype is None:
        np_dtype = like.dtype
    else:
        np_dtype = {'int64': np.int64, 'float32': np.float32, 'bool': bool}[dtype]
    return np.empty(size, dtype=np_dtype)


def empty_like(arr: Array) -> Array:
    """
    Create an empty array with same shape/dtype/device as input.

    Args:
        arr: Template array.

    Returns:
        Empty array with same properties as input.
    """
    if is_torch(arr):
        import torch
        return torch.empty_like(arr)
    return np.empty_like(arr)


def clone(arr: Array, detach: bool = True) -> Array:
    """
    Create a copy of an array, optionally detached from any computation graph.

    For PyTorch tensors, this returns a clone (detached by default).
    For NumPy arrays, this returns a copy.

    Args:
        arr: Array to clone.
        detach: If True (default), detach from computation graph.
            Set to False to allow gradients to flow through the clone.

    Returns:
        Copy of the array.
    """
    if is_torch(arr):
        if detach:
            return arr.detach().clone()
        return arr.clone()
    return arr.copy()


# =============================================================================
# Utility Operations
# =============================================================================

def nonzero_1d(arr: Array) -> Array:
    """
    Get indices of non-zero elements in a 1D array.

    Args:
        arr: 1D array.

    Returns:
        Indices of non-zero elements in original backend.
    """
    if is_torch(arr):
        return arr.nonzero().squeeze(-1)
    return arr.nonzero()[0]


def to_int64(arr: Array) -> Array:
    """
    Convert array to int64 dtype.

    Args:
        arr: Input array.

    Returns:
        Array with int64 dtype in original backend.
    """
    if is_torch(arr):
        return arr.long()
    return arr.astype(np.int64)


def to_float64(arr: Array) -> Array:
    """
    Convert array to float64 dtype.

    Args:
        arr: Input array.

    Returns:
        Array with float64 dtype in original backend.
    """
    if is_torch(arr):
        return arr.double()
    return arr.astype(np.float64)


def to_dtype_of(arr: Array, like: Array) -> Array:
    """
    Convert array to match the dtype of another array.

    Args:
        arr: Input array to convert.
        like: Reference array whose dtype to match.

    Returns:
        Array with same dtype as `like`.
    """
    if is_torch(arr):
        return arr.to(like.dtype)
    return arr.astype(like.dtype)


def isin(arr: Array, values: Array) -> Array:
    """
    Check if elements of arr are in values.

    Args:
        arr: Input array to check.
        values: Array of values to check against (must match arr backend).

    Returns:
        Boolean array of same shape as arr, True where element is in values.

    Raises:
        TypeError: If arr and values have different backends.
    """
    _check_backends(arr, values)
    if is_torch(arr):
        import torch
        return torch.isin(arr, values.to(device=arr.device, dtype=arr.dtype))
    return np.isin(arr, values)


# =============================================================================
# Math Operations
# =============================================================================

def sqrt(arr: Array) -> Array:
    """
    Element-wise square root.

    Args:
        arr: Input array.

    Returns:
        Square root in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.sqrt(arr)
    return np.sqrt(arr)


def sin(arr: Array) -> Array:
    """
    Element-wise sine.

    Args:
        arr: Input array (angles in radians).

    Returns:
        Sine values in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.sin(arr)
    return np.sin(arr)


def cos(arr: Array) -> Array:
    """
    Element-wise cosine.

    Args:
        arr: Input array (angles in radians).

    Returns:
        Cosine values in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.cos(arr)
    return np.cos(arr)


def acos(arr: Array) -> Array:
    """
    Element-wise arc cosine (inverse cosine).

    Args:
        arr: Input array with values in [-1, 1].

    Returns:
        Arc cosine values in radians, in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.acos(arr)
    return np.arccos(arr)


def log(arr: Array) -> Array:
    """
    Element-wise natural logarithm.

    Args:
        arr: Input array.

    Returns:
        Log values in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.log(arr)
    return np.log(arr)


def exp(arr: Array) -> Array:
    """
    Element-wise exponential.

    Args:
        arr: Input array.

    Returns:
        Exponential values in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.exp(arr)
    return np.exp(arr)


def norm(arr: Array, axis: int | None = None, keepdims: bool = False) -> Array:
    """
    Compute vector norm along an axis.

    Args:
        arr: Input array.
        axis: Axis along which to compute norm (None for full array).
        keepdims: If True, keep the reduced dimension.

    Returns:
        Norm values in original backend.
    """
    if is_torch(arr):
        import torch
        if axis is None:
            return torch.norm(arr)
        return torch.norm(arr, dim=axis, keepdim=keepdims)
    return np.linalg.norm(arr, axis=axis, keepdims=keepdims)


def sigmoid(arr: Array) -> Array:
    """
    Element-wise sigmoid (logistic) function.

    Args:
        arr: Input array.

    Returns:
        Sigmoid values in [0, 1].
    """
    if is_torch(arr):
        import torch
        return torch.sigmoid(arr)
    # Clip to avoid overflow in exp (exp(88) ≈ 1.6e38 < float32 max)
    arr = np.clip(arr, -88, 88)
    return 1.0 / (1.0 + np.exp(-arr))


def prod(arr: Array, axis: int | None = None) -> Array:
    """
    Product of elements along an axis.

    Args:
        arr: Input array.
        axis: Axis along which to compute product. None for all elements.

    Returns:
        Product values in original backend.
    """
    if is_torch(arr):
        import torch
        if axis is None:
            return torch.prod(arr)
        return torch.prod(arr, dim=axis)
    return np.prod(arr, axis=axis)


def mean(arr: Array, axis: int | None = None) -> Array:
    """
    Mean of elements along an axis.

    Args:
        arr: Input array.
        axis: Axis along which to compute mean. None for all elements.

    Returns:
        Mean values in original backend.
    """
    if is_torch(arr):
        import torch
        if axis is None:
            return torch.mean(arr)
        return torch.mean(arr, dim=axis)
    return np.mean(arr, axis=axis)


def eye(n: int, *, like: Array) -> Array:
    """
    Create an identity matrix matching the backend/device of 'like'.

    Args:
        n: Size of the identity matrix (n x n).
        like: Template array for backend and device.

    Returns:
        Identity matrix in the same backend/dtype/device as 'like'.
    """
    if is_torch(like):
        import torch
        return torch.eye(n, device=like.device, dtype=like.dtype)
    return np.eye(n, dtype=like.dtype)


def zeros_like(arr: Array) -> Array:
    """
    Create a zeros array with same shape/dtype/device as input.

    Args:
        arr: Template array.

    Returns:
        Zeros array with same properties as input.
    """
    if is_torch(arr):
        import torch
        return torch.zeros_like(arr)
    return np.zeros_like(arr)


def fill_diagonal(arr: Array, value: float) -> Array:
    """
    Fill the diagonal of a 2D array with a value (in-place).

    Args:
        arr: 2D array to modify.
        value: Value to fill on the diagonal.

    Returns:
        The modified array (same object, modified in-place).
    """
    if is_torch(arr):
        arr.fill_diagonal_(value)
    else:
        np.fill_diagonal(arr, value)
    return arr


def zeros_nd(shape: tuple, *, like: Array) -> Array:
    """
    Create a zeros array with specified shape, matching dtype/device of 'like'.

    Args:
        shape: Shape of the output array.
        like: Template array for dtype and device.

    Returns:
        Zeros array with specified shape.
    """
    if is_torch(like):
        import torch
        return torch.zeros(shape, dtype=like.dtype, device=like.device)
    return np.zeros(shape, dtype=like.dtype)


def ones_like(arr: Array) -> Array:
    """
    Create a ones array with same shape/dtype/device as input.

    Args:
        arr: Template array.

    Returns:
        Ones array with same properties as input.
    """
    if is_torch(arr):
        import torch
        return torch.ones_like(arr)
    return np.ones_like(arr)


def bmm(a: Array, b: Array) -> Array:
    """
    Batched matrix multiplication.

    Args:
        a: First batch of matrices, shape (batch, n, m).
        b: Second batch of matrices, shape (batch, m, p).

    Returns:
        Batched product, shape (batch, n, p).
    """
    _check_backends(a, b)
    if is_torch(a):
        import torch
        return torch.bmm(a, b)
    return np.einsum('bij,bjk->bik', a, b)


def where(condition: Array, x: Array, y: Array) -> Array:
    """
    Select elements from x or y based on condition.

    Args:
        condition: Boolean array.
        x: Values where condition is True.
        y: Values where condition is False.

    Returns:
        Array with selected values.
    """
    if is_torch(condition):
        import torch
        return torch.where(condition, x, y)
    return np.where(condition, x, y)


def unsqueeze(arr: Array, axis: int) -> Array:
    """
    Add a dimension of size 1 at the specified axis.

    Args:
        arr: Input array.
        axis: Position where the new axis should be inserted.

    Returns:
        Array with an additional dimension.
    """
    if is_torch(arr):
        return arr.unsqueeze(axis)
    return np.expand_dims(arr, axis=axis)


def expand(arr: Array, shape: tuple) -> Array:
    """
    Broadcast array to a larger shape.

    Uses -1 in shape to keep existing dimension size.

    Args:
        arr: Input array.
        shape: Target shape with -1 for unchanged dimensions.

    Returns:
        Broadcasted array (may be a view).
    """
    if is_torch(arr):
        return arr.expand(*shape)
    # NumPy: convert -1 to actual sizes, then broadcast
    actual_shape = tuple(
        arr.shape[i] if s == -1 else s
        for i, s in enumerate(shape)
    )
    return np.broadcast_to(arr, actual_shape)


def transpose(arr: Array, axes: tuple) -> Array:
    """
    Permute array dimensions.

    Args:
        arr: Input array.
        axes: Tuple specifying the new order of dimensions.

    Returns:
        Array with permuted dimensions.
    """
    if is_torch(arr):
        return arr.permute(*axes)
    return np.transpose(arr, axes)


def outer(a: Array, b: Array) -> Array:
    """
    Outer product of two 1D arrays.

    Args:
        a: First 1D array of shape (M,).
        b: Second 1D array of shape (N,).

    Returns:
        Outer product matrix of shape (M, N).
    """
    _check_backends(a, b)
    if is_torch(a):
        import torch
        return torch.outer(a, b)
    return np.outer(a, b)


def clamp(arr: Array, min_val: float | None = None, max_val: float | None = None) -> Array:
    """
    Clamp array values to a range.

    Args:
        arr: Input array.
        min_val: Minimum value (None for no lower bound).
        max_val: Maximum value (None for no upper bound).

    Returns:
        Clamped array in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.clamp(arr, min=min_val, max=max_val)
    result = arr
    if min_val is not None:
        result = np.maximum(result, min_val)
    if max_val is not None:
        result = np.minimum(result, max_val)
    return result


def argsort(arr: Array) -> Array:
    """
    Return indices that would sort the array.

    Args:
        arr: 1D input array.

    Returns:
        Array of indices that would sort the array.
    """
    if is_torch(arr):
        import torch
        return torch.argsort(arr)
    return np.argsort(arr)


def diff(arr: Array) -> Array:
    """
    Compute differences between consecutive elements.

    Args:
        arr: 1D input array.

    Returns:
        Array of differences (length n-1 for input of length n).
    """
    if is_torch(arr):
        import torch
        return torch.diff(arr)
    return np.diff(arr)


def split_at_indices(arr: Array, split_indices: Array) -> list:
    """
    Split array at specified indices.

    Args:
        arr: Array to split.
        split_indices: 1D array of indices where splits occur.

    Returns:
        List of array chunks.
    """
    if is_torch(arr):
        import torch
        # Convert split_indices to list for torch.tensor_split
        if len(split_indices) == 0:
            return [arr]
        return list(torch.tensor_split(arr, split_indices.cpu().tolist()))
    # NumPy path
    if len(split_indices) == 0:
        return [arr]
    return list(np.split(arr, split_indices))


def topk(arr: Array, k: int, dim: int = -1, largest: bool = True) -> tuple[Array, Array]:
    """
    Find k largest or smallest elements along a dimension.

    Args:
        arr: Input array.
        k: Number of elements to return.
        dim: Dimension along which to find topk.
        largest: If True, find k largest; if False, find k smallest.

    Returns:
        Tuple of (values, indices) in original backend.
    """
    if is_torch(arr):
        import torch
        return torch.topk(arr, k, dim=dim, largest=largest)
    else:
        # NumPy implementation
        if largest:
            # Get indices of k largest
            indices = np.argpartition(arr, -k, axis=dim)
            indices = np.take(indices, range(-k, 0), axis=dim)
        else:
            # Get indices of k smallest
            indices = np.argpartition(arr, k, axis=dim)
            indices = np.take(indices, range(k), axis=dim)

        # Get values and sort by value
        values = np.take_along_axis(arr, indices, axis=dim)
        sort_idx = np.argsort(values, axis=dim)
        if largest:
            sort_idx = np.flip(sort_idx, axis=dim)
        indices = np.take_along_axis(indices, sort_idx, axis=dim)
        values = np.take_along_axis(values, sort_idx, axis=dim)

        return values, indices
