# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

import inspect
import math
from functools import lru_cache

from awkward._nplikes.array_like import maybe_materialize
from awkward._nplikes.numpy_like import (
    ArrayLike,
    IndexType,
    NumpyLike,
    NumpyMetadata,
    UfuncLike,
    UniqueAllResult,
)
from awkward._nplikes.placeholder import PlaceholderArray
from awkward._nplikes.shape import ShapeItem, unknown_length
from awkward._nplikes.virtual import VirtualNDArray
from awkward._typing import TYPE_CHECKING, Any, DType, Final, Literal, TypeVar, cast

if TYPE_CHECKING:
    from numpy.typing import DTypeLike

np = NumpyMetadata.instance()


@lru_cache
def _nplike_concatenate_has_casting(module: Any) -> bool:
    x = module.zeros(2)
    try:
        module.concatenate((x, x), casting="same_kind")
    except TypeError:
        return False
    else:
        return True


@lru_cache
def _nplike_reshape_has_copy(module: Any) -> bool:
    x = module.zeros(2)
    try:
        module.reshape(x, (2,), copy=False)
    except TypeError:
        return False
    else:
        return True


ArrayLikeT = TypeVar("ArrayLikeT", bound=ArrayLike)


class ArrayModuleNumpyLike(NumpyLike[ArrayLikeT]):
    """
    Some methods maintain virtualness while others are required to materialize the array.
    The `maybe_materialize` function is used for that purpose.
    If the input is not a virtual array, it's a zero cost operation.
    """

    known_data: Final[bool] = True
    _module: Any

    def prepare_ufunc(self, ufunc: UfuncLike) -> UfuncLike:
        return ufunc

    ############################ array creation

    def asarray(
        self,
        obj,
        *,
        dtype: DTypeLike | None = None,
        copy: bool | None = None,
    ) -> ArrayLikeT | PlaceholderArray | VirtualNDArray:
        if isinstance(obj, PlaceholderArray):
            assert obj.dtype == dtype or dtype is None
            return obj
        if isinstance(obj, VirtualNDArray):
            if obj.is_materialized:
                return self.asarray(obj.materialize(), dtype=dtype, copy=copy)
            else:
                # if we are not copying and the dtype is _exactly_ the dtype of the existing array
                # or dtype is None, we can return the VirtualNDArray directly
                # this avoids unnecessary VirtualNDArray creation and method-chaining
                if not copy and (obj.dtype == dtype or dtype is None):
                    return obj
                return VirtualNDArray(
                    obj._nplike,
                    obj._shape,
                    obj._dtype if dtype is None else dtype,
                    lambda: self.asarray(obj.materialize(), dtype=dtype, copy=copy),
                    lambda: obj.shape,
                    obj._buffer_key,
                    __enable_caching__=obj.__enable_caching__,
                )
        if copy:
            return self._module.array(obj, dtype=dtype, copy=True)
        elif copy is None:
            return self._module.asarray(obj, dtype=dtype)
        else:
            if dtype is not None and getattr(obj, "dtype", dtype) != dtype:
                raise ValueError(
                    "asarray was called with copy=False for an array of a different dtype"
                )
            else:
                return self._module.asarray(obj, dtype=dtype)

    def ascontiguousarray(
        self, x: ArrayLikeT | PlaceholderArray
    ) -> ArrayLikeT | PlaceholderArray | VirtualNDArray:
        if isinstance(x, PlaceholderArray):
            return x
        elif isinstance(x, VirtualNDArray):
            if x.is_materialized:
                return self.ascontiguousarray(x.materialize())  #  type: ignore[arg-type]
            else:
                return VirtualNDArray(
                    x._nplike,
                    x._shape,
                    x._dtype,
                    lambda: self.ascontiguousarray(x.materialize()),  # type: ignore[arg-type]
                    lambda: x.shape,
                    x._buffer_key,
                    __enable_caching__=x.__enable_caching__,
                )
        else:
            return self._module.ascontiguousarray(x)

    def frombuffer(
        self, buffer, *, dtype: DTypeLike | None = None, count: ShapeItem = -1
    ) -> ArrayLikeT:
        if isinstance(buffer, PlaceholderArray):
            raise TypeError("placeholder arrays are not supported in `frombuffer`")
        if isinstance(buffer, VirtualNDArray):
            raise TypeError("virtual arrays are not supported in `frombuffer`")
        return self._module.frombuffer(buffer, dtype=dtype, count=count)

    def from_dlpack(self, x: Any) -> ArrayLikeT:
        return self._module.from_dlpack(x)

    def zeros(
        self,
        shape: ShapeItem | tuple[ShapeItem, ...],
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        return self._module.zeros(shape, dtype=dtype)

    def ones(
        self,
        shape: ShapeItem | tuple[ShapeItem, ...],
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        return self._module.ones(shape, dtype=dtype)

    def empty(
        self,
        shape: ShapeItem | tuple[ShapeItem, ...],
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        return self._module.empty(shape, dtype=dtype)

    def full(
        self,
        shape: ShapeItem | tuple[ShapeItem, ...],
        fill_value,
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        return self._module.full(shape, self._module.array(fill_value), dtype=dtype)

    def zeros_like(
        self, x: ArrayLikeT | PlaceholderArray, *, dtype: DTypeLike | None = None
    ) -> ArrayLikeT:
        if isinstance(x, (PlaceholderArray, VirtualNDArray)):
            return self.zeros(x.shape, dtype=dtype or x.dtype)
        else:
            return self._module.zeros_like(x, dtype=dtype)

    def ones_like(
        self, x: ArrayLikeT | PlaceholderArray, *, dtype: DTypeLike | None = None
    ) -> ArrayLikeT:
        if isinstance(x, (PlaceholderArray, VirtualNDArray)):
            return self.ones(x.shape, dtype=dtype or x.dtype)
        else:
            return self._module.ones_like(x, dtype=dtype)

    def full_like(
        self,
        x: ArrayLikeT | PlaceholderArray,
        fill_value,
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        if isinstance(x, (PlaceholderArray, VirtualNDArray)):
            return self.full(x.shape, fill_value, dtype=dtype or x.dtype)
        else:
            return self._module.full_like(
                x, self._module.array(fill_value), dtype=dtype
            )

    def arange(
        self,
        start: float | int,
        stop: float | int | None = None,
        step: float | int = 1,
        *,
        dtype: DTypeLike | None = None,
    ) -> ArrayLikeT:
        start, stop, step = maybe_materialize(start, stop, step)
        return self._module.arange(start, stop, step, dtype=dtype)

    def meshgrid(
        self, *arrays: ArrayLikeT, indexing: Literal["xy", "ij"] = "xy"
    ) -> list[ArrayLikeT]:
        return self._module.meshgrid(*maybe_materialize(*arrays), indexing=indexing)

    ############################ testing

    def array_equal(
        self, x1: ArrayLikeT, x2: ArrayLikeT, *, equal_nan: bool = False
    ) -> bool:
        x1, x2 = maybe_materialize(x1, x2)
        if equal_nan:
            # Only newer numpy.array_equal supports the equal_nan parameter.
            both_nan = self._module.logical_and(
                self._module.isnan(x1), self._module.isnan(x2)
            )
            both_equal = x1 == x2
            return self._module.all(self._module.logical_or(both_equal, both_nan))
        else:
            return self._module.array_equal(x1, x2)

    def searchsorted(
        self,
        x: ArrayLikeT,
        values: ArrayLikeT,
        *,
        side: Literal["left", "right"] = "left",
        sorter: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        x, values, sorter = maybe_materialize(x, values, sorter)
        return self._module.searchsorted(x, values, side=side, sorter=sorter)

    ############################ manipulation
    def apply_ufunc(
        self,
        ufunc: UfuncLike,
        method: str,
        args: list[Any],
        kwargs: dict[str, Any] | None = None,
    ) -> ArrayLikeT | tuple[ArrayLikeT, ...]:
        if method != "__call__" or len(args) == 0:
            raise NotImplementedError

        args = list(maybe_materialize(*args))

        if hasattr(ufunc, "resolve_dtypes"):
            return self._apply_ufunc_nep_50(ufunc, method, args, kwargs)
        else:
            return self._apply_ufunc_legacy(ufunc, method, args, kwargs)

    def _get_nep_50_dtype(
        self, obj: Any
    ) -> DType | type[int] | type[complex] | type[float]:
        if hasattr(obj, "dtype"):
            return obj.dtype
        elif isinstance(obj, bool):
            return np.dtype(np.bool_)
        else:
            assert isinstance(obj, (int, complex, float))
            return type(obj)

    # Does NumPy support value-less ufunc resolution?
    def _apply_ufunc_nep_50(
        self,
        ufunc: UfuncLike,
        method: str,
        args: list[Any],
        kwargs: dict[str, Any] | None = None,
    ) -> ArrayLikeT | tuple[ArrayLikeT]:
        # Determine input argument dtypes
        input_arg_dtypes = [self._get_nep_50_dtype(obj) for obj in args]
        # Resolve these for the given ufunc
        arg_dtypes = tuple(input_arg_dtypes + [None] * ufunc.nout)
        resolved_dtypes = ufunc.resolve_dtypes(arg_dtypes)
        # Interpret the arguments under these dtypes, converting scalars to length-1 arrays
        resolved_args = [
            cast("ArrayLikeT", self.asarray(arg, dtype=dtype))
            for arg, dtype in zip(args, resolved_dtypes[: len(args)], strict=True)
        ]
        # Broadcast to ensure all-scalar or all-nd-array
        broadcasted_args = self.broadcast_arrays(*resolved_args)
        # Allow other nplikes to replace implementation
        impl = self.prepare_ufunc(ufunc)
        # Compute the result
        return impl(*broadcasted_args, **(kwargs or {}))

    # Otherwise, perform default NumPy coercion (value-dependent)
    def _apply_ufunc_legacy(
        self,
        ufunc: UfuncLike,
        method: str,
        args: list[Any],
        kwargs: dict[str, Any] | None = None,
    ) -> ArrayLikeT | tuple[ArrayLikeT]:
        # Convert np.generic to scalar arrays
        resolved_args = [
            cast(
                "ArrayLikeT",
                self.asarray(arg, dtype=arg.dtype if hasattr(arg, "dtype") else None),
            )
            for arg in args
        ]
        broadcasted_args = self.broadcast_arrays(*resolved_args)
        # Choose the broadcasted argument if it wasn't a Python scalar
        non_generic_value_promoted_args = [
            y if hasattr(x, "ndim") else x
            for x, y in zip(args, broadcasted_args, strict=True)
        ]
        # Allow other nplikes to replace implementation
        impl = self.prepare_ufunc(ufunc)
        # Compute the result
        return impl(*non_generic_value_promoted_args, **(kwargs or {}))

    def broadcast_arrays(self, *arrays: ArrayLikeT) -> list[ArrayLikeT]:
        arrays = maybe_materialize(*arrays)
        return self._module.broadcast_arrays(*arrays)

    def _compute_compatible_shape(
        self, shape: tuple[ShapeItem, ...], existing_shape: tuple[ShapeItem, ...]
    ) -> tuple[ShapeItem, ...]:
        next_shape = list(shape)
        j = None
        length_factor: ShapeItem = 1
        for i, item in enumerate(shape):
            if item != -1:
                length_factor *= item
            elif j is not None:
                raise ValueError("can only have one unknown dimension")
            else:
                j = i
        if j is not None:
            next_shape[j] = math.prod(existing_shape) // length_factor
        return tuple(next_shape)

    def reshape(
        self,
        x: ArrayLikeT | PlaceholderArray,
        shape: tuple[ShapeItem, ...],
        *,
        copy: bool | None = None,
    ) -> ArrayLikeT | PlaceholderArray | VirtualNDArray:
        if isinstance(x, PlaceholderArray):
            next_shape = self._compute_compatible_shape(shape, x.shape)
            return PlaceholderArray(self, next_shape, x.dtype, x._buffer_key)
        if isinstance(x, VirtualNDArray):
            if x.is_materialized:
                return self.reshape(x.materialize(), shape, copy=copy)  # type: ignore[arg-type]
            else:
                next_shape = self._compute_compatible_shape(shape, x.shape)
                # if the reshape is _exactly_ shaping the array as it is already, we can return the VirtualNDArray directly
                # this avoids unnecessary VirtualNDArray creation and method-chaining
                if next_shape == x.shape and not copy:
                    return x
                return VirtualNDArray(
                    self,
                    next_shape,
                    x.dtype,
                    lambda: self.reshape(x.materialize(), next_shape, copy=copy),  # type: ignore[arg-type]
                    None,
                    x._buffer_key,
                    __enable_caching__=x.__enable_caching__,
                )

        if copy is None:
            return self._module.reshape(x, shape)
        elif copy:
            return self._module.reshape(self._module.copy(x, order="C"), shape)
        else:
            result = self._module.asarray(x)
            if _nplike_reshape_has_copy(self._module):
                return self._module.reshape(result, shape, copy=False)
            try:
                result.shape = shape
            except AttributeError:
                raise ValueError(
                    "Unable to avoid creating a copy while reshaping."
                ) from None
            return result

    def shape_item_as_index(self, x1: ShapeItem) -> int:
        if x1 is unknown_length:
            raise TypeError("array module nplikes do not support unknown lengths")
        elif isinstance(x1, int):
            return x1
        else:
            raise TypeError(f"expected None or int type, received {x1}")

    def index_as_shape_item(self, x1: IndexType) -> int:
        return int(x1)

    def derive_slice_for_length(
        self, slice_: slice, length: ShapeItem
    ) -> tuple[IndexType, IndexType, IndexType, ShapeItem]:
        """
        Args:
            slice_: normalized slice object
            length: length of layout

        Return a tuple of (start, stop, step) indices into a layout, suitable for
        `_getitem_range` (if step == 1). Normalize lengths to fit length of array,
        and for arrays with unknown lengths, these offsets become none.
        """
        # We have known_data (therefore known shape), so we can safely convert to int
        slice_as_shape = slice(
            slice_.start
            if slice_.start is None
            else self.index_as_shape_item(slice_.start),
            slice_.stop
            if slice_.stop is None
            else self.index_as_shape_item(slice_.stop),
            slice_.step
            if slice_.step is None
            else self.index_as_shape_item(slice_.step),
        )
        start, stop, step = slice_as_shape.indices(length)
        slice_length = math.ceil((stop - start) / step)

        # Shape items are already valid indices, so we don't need to invoke
        # `shape_item_as_index`
        return start, stop, step, slice_length

    def regularize_index_for_length(
        self, index: IndexType, length: ShapeItem
    ) -> IndexType:
        """
        Args:
            index: index value
            length: length of array

        Returns regularized index that is guaranteed to be in-bounds.
        """  # We have known length and index
        length = cast(int, length)

        if index < 0:
            index = index + length

        if 0 <= index < length:
            return index
        else:
            raise IndexError(f"index value out of bounds (0, {length}): {index}")

    def nonzero(self, x: ArrayLikeT) -> tuple[ArrayLikeT, ...]:
        (x,) = maybe_materialize(x)
        return self._module.nonzero(x)

    def where(
        self, condition: ArrayLikeT, x1: ArrayLikeT, x2: ArrayLikeT
    ) -> ArrayLikeT:
        condition, x1, x2 = maybe_materialize(condition, x1, x2)

        return self._module.where(condition, x1, x2)

    def unique_values(self, x: ArrayLikeT) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        np_unique_accepts_equal_nan = (
            "equal_nan" in inspect.signature(self._module.unique).parameters
        )

        if np_unique_accepts_equal_nan:
            return self._module.unique(
                x,
                return_counts=False,
                return_index=False,
                return_inverse=False,
                equal_nan=False,
            )
        else:
            return self._module.unique(
                x,
                return_counts=False,
                return_index=False,
                return_inverse=False,
            )

    def unique_all(self, x: ArrayLikeT) -> UniqueAllResult:
        (x,) = maybe_materialize(x)
        values, indices, inverse_indices, counts = self._module.unique(
            x, return_counts=True, return_index=True, return_inverse=True
        )
        # np.unique() flattens inverse indices, but they need to share x's shape
        # See https://github.com/numpy/numpy/issues/20638
        inverse_indices = inverse_indices.reshape(x.shape)
        return UniqueAllResult(values, indices, inverse_indices, counts)

    def sort(
        self,
        x: ArrayLikeT,
        *,
        axis: int = -1,
        descending: bool = False,
        stable: bool = True,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        # Note: this keyword argument is different, and the default is different.
        kind = "stable" if stable else "quicksort"
        res = self._module.sort(x, axis=axis, kind=kind)
        if descending:
            return self._module.flip(res, axis=axis)
        else:
            return res

    def concat(
        self,
        arrays: list[ArrayLikeT] | tuple[ArrayLikeT, ...],
        *,
        axis: int | None = 0,
    ) -> ArrayLikeT | VirtualNDArray:
        if any(isinstance(x, VirtualNDArray) and not x.is_materialized for x in arrays):
            # TODO: switch to np.result_type when we pin numpy >= 1.23
            dtype = self._module.concatenate(
                [self._module.empty(0, x.dtype) for x in arrays]
            ).dtype

            if axis is None:
                shape = (sum((x.size for x in arrays), 0),)
            else:
                ref = arrays[0]
                if axis < 0:
                    axis += ref.ndim
                for i, x in enumerate(arrays[1:], start=1):
                    for dim, (ref_d, x_d) in enumerate(
                        zip(ref.shape, x.shape, strict=True)
                    ):
                        if dim != axis and ref_d != x_d:
                            raise ValueError(
                                "all the input array dimensions except for the concatenation axis "
                                f"must match exactly, but along dimension {dim}, the array at index 0 "
                                f"has size {ref_d} and the array at index {i} has size {x_d}"
                            )
                shape = (  # type: ignore[assignment]
                    *ref.shape[:axis],
                    sum((x.shape[axis] for x in arrays), 0),
                    *ref.shape[axis + 1 :],
                )

            return VirtualNDArray(
                self,
                shape,
                dtype,
                lambda: self.concat(maybe_materialize(*arrays), axis=axis),
            )

        arrays = maybe_materialize(*arrays)
        if _nplike_concatenate_has_casting(self._module):
            return self._module.concatenate(arrays, axis=axis, casting="same_kind")
        else:
            return self._module.concatenate(arrays, axis=axis)

    def repeat(
        self,
        x: ArrayLikeT,
        repeats: ArrayLikeT | int,
        *,
        axis: int | None = None,
    ) -> ArrayLikeT:
        x, repeats = maybe_materialize(x, repeats)
        return self._module.repeat(x, repeats=repeats, axis=axis)

    def stack(
        self,
        arrays: list[ArrayLikeT] | tuple[ArrayLikeT, ...],
        *,
        axis: int = 0,
    ) -> ArrayLikeT:
        arrays = maybe_materialize(*arrays)
        arrays = list(arrays)
        return self._module.stack(arrays, axis=axis)

    def packbits(
        self,
        x: ArrayLikeT,
        *,
        axis: int | None = None,
        bitorder: Literal["big", "little"] = "big",
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.packbits(x, axis=axis, bitorder=bitorder)

    def unpackbits(
        self,
        x: ArrayLikeT,
        *,
        axis: int | None = None,
        count: int | None = None,
        bitorder: Literal["big", "little"] = "big",
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.unpackbits(x, axis=axis, count=count, bitorder=bitorder)

    def broadcast_to(self, x: ArrayLikeT, shape: tuple[ShapeItem, ...]) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.broadcast_to(x, shape)

    def strides(self, x: ArrayLikeT | PlaceholderArray) -> tuple[ShapeItem, ...]:
        if isinstance(x, PlaceholderArray):
            # Assume contiguous
            strides: tuple[ShapeItem, ...] = (x.dtype.itemsize,)
            for item in x.shape[-1:0:-1]:
                strides = (item * strides[0], *strides)
            return strides

        (x,) = maybe_materialize(x)
        return x.strides

    def byteswap(self, x: ArrayLikeT) -> ArrayLikeT:
        raise NotImplementedError

    ############################ ufuncs

    def add(
        self, x1: ArrayLikeT, x2: ArrayLikeT, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.add(x1, x2, out=maybe_out)

    def logical_or(
        self, x1: ArrayLikeT, x2: ArrayLikeT, *, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.logical_or(x1, x2, out=maybe_out)

    def logical_and(
        self, x1: ArrayLikeT, x2: ArrayLikeT, *, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.logical_and(x1, x2, out=maybe_out)

    def logical_not(
        self, x: ArrayLikeT, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.logical_not(x, out=maybe_out)

    def sqrt(self, x: ArrayLikeT, maybe_out: ArrayLikeT | None = None) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.sqrt(x, out=maybe_out)

    def exp(self, x: ArrayLikeT, maybe_out: ArrayLikeT | None = None) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.exp(x, out=maybe_out)

    def divide(
        self, x1: ArrayLikeT, x2: ArrayLikeT, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.divide(x1, x2, out=maybe_out)

    def minimum(
        self, x1: ArrayLikeT, x2: ArrayLikeT, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.minimum(x1, x2, out=maybe_out)

    def maximum(
        self, x1: ArrayLikeT, x2: ArrayLikeT, maybe_out: ArrayLikeT | None = None
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.maximum(x1, x2, out=maybe_out)

    ############################ almost-ufuncs

    def nan_to_num(
        self,
        x: ArrayLikeT,
        *,
        copy: bool = True,
        nan: int | float | None = 0.0,
        posinf: int | float | None = None,
        neginf: int | float | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.nan_to_num(
            x, copy=copy, nan=nan, posinf=posinf, neginf=neginf
        )

    def isclose(
        self,
        x1: ArrayLikeT,
        x2: ArrayLikeT,
        *,
        rtol: float = 1e-5,
        atol: float = 1e-8,
        equal_nan: bool = False,
    ) -> ArrayLikeT:
        x1, x2 = maybe_materialize(x1, x2)
        return self._module.isclose(x1, x2, rtol=rtol, atol=atol, equal_nan=equal_nan)

    def isnan(self, x: ArrayLikeT) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.isnan(x)

    def all(
        self,
        x: ArrayLikeT,
        *,
        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
        keepdims: bool = False,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.all(x, axis=axis, keepdims=keepdims, out=maybe_out)

    def any(
        self,
        x: ArrayLikeT,
        *,
        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
        keepdims: bool = False,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.any(x, axis=axis, keepdims=keepdims, out=maybe_out)

    def min(
        self,
        x: ArrayLikeT,
        *,
        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
        keepdims: bool = False,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.min(x, axis=axis, keepdims=keepdims, out=maybe_out)

    def sum(
        self,
        x: ArrayLikeT,
        *,
        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
        keepdims: bool = False,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.sum(x, axis=axis, keepdims=keepdims, out=maybe_out)

    def max(
        self,
        x: ArrayLikeT,
        *,
        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
        keepdims: bool = False,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.max(x, axis=axis, keepdims=keepdims, out=maybe_out)

    def count_nonzero(
        self, x: ArrayLikeT, *, axis: ShapeItem | tuple[ShapeItem, ...] | None = None
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        assert isinstance(axis, int) or axis is None
        return self._module.count_nonzero(x, axis=axis)

    def cumsum(
        self,
        x: ArrayLikeT,
        *,
        axis: int | None = None,
        maybe_out: ArrayLikeT | None = None,
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.cumsum(x, axis=axis, out=maybe_out)

    def real(self, x: ArrayLikeT) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        xr = self._module.real(x)
        # For numpy, xr is a view on x, but we don't want to mutate x.
        return self._module.copy(xr)

    def imag(self, x: ArrayLikeT) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        xr = self._module.imag(x)
        # For numpy, xr is a view on x, but we don't want to mutate x.
        return self._module.copy(xr)

    def angle(self, x: ArrayLikeT, deg: bool = False) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.angle(x, deg)

    def round(self, x: ArrayLikeT, decimals: int = 0) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return self._module.round(x, decimals=decimals)

    def array_str(
        self,
        x: ArrayLikeT,
        *,
        max_line_width: int | None = None,
        precision: int | None = None,
        suppress_small: bool | None = None,
    ):
        if isinstance(x, VirtualNDArray) and not x.is_materialized:
            return "[## ... ##]"
        if isinstance(x, PlaceholderArray):
            return "[## ... ##]"
        (x,) = maybe_materialize(x)
        return self._module.array_str(
            x,
            max_line_width=max_line_width,
            precision=precision,
            suppress_small=suppress_small,
        )

    def astype(
        self, x: ArrayLikeT, dtype: DTypeLike, *, copy: bool | None = True
    ) -> ArrayLikeT:
        (x,) = maybe_materialize(x)
        return x.astype(dtype, copy=copy)  # type: ignore[attr-defined]

    def can_cast(
        self,
        from_: DTypeLike | ArrayLikeT,
        to: DTypeLike | ArrayLikeT,
        casting: Literal["no", "equiv", "safe", "same_kind", "unsafe"] = "same_kind",
    ) -> bool:
        return self._module.can_cast(from_, to, casting=casting)

    @classmethod
    def is_own_array(cls, obj) -> bool:
        # This kinda breaks the type(obj) -> ownership paradigm, but I don't see a simpler way for virtual arrays.
        # We are required to know if different nplikes own virtual arrays throughout the code
        # and that can only be determined by the underlying nplike of the virtual array
        # as virtual arrays can generate either numpy or cupy ndarrays.
        # The VirtualNDArray type is now enough. We need the underlying nplike to determine ownership.
        # All nplikes implement an ndarray property so we can can use that.
        # There is an extra isinstance check here but that has to live somewhere in the code either way to determine ownership.
        if isinstance(obj, VirtualNDArray):
            return cls.is_own_array_type(obj.nplike.ndarray)
        return cls.is_own_array_type(type(obj))

    def memory_ptr(self, x: ArrayLikeT) -> int:
        raise NotImplementedError
