#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
This file includes utility functions for arrays/tensors that work for multiple
libraries like NumPy and PyTorch.
"""

import builtins
import functools
import math
import numbers

from polygraphy import mod
from polygraphy.datatype import DataType
from polygraphy.logger import G_LOGGER

np = mod.lazy_import("numpy")
torch = mod.lazy_import("torch>=1.13.0")


@mod.export()
def is_torch(obj):
    """
    Whether the provided object is a PyTorch tensor.
    This function does *not* introduce a dependency on the PyTorch module.

    Args:
        obj (Any): The object to check.

    Returns:
        bool: Whether the object is a PyTorch tensor.
    """
    return (
        torch.is_installed() and torch.is_importable() and isinstance(obj, torch.Tensor)
    )


@mod.export()
def is_numpy(obj):
    """
    Whether the provided object is a NumPy array or scalar.
    This function does *not* introduce a dependency on the NumPy module.

    Args:
        obj (Any): The object to check.

    Returns:
        bool: Whether the object is a NumPy array.
    """
    return (
        np.is_installed()
        and np.is_importable()
        and (isinstance(obj, np.ndarray) or isinstance(obj, np.generic))
    )


@mod.export()
def is_device_view(obj):
    """
    Whether the provided object is a DeviceView array.

    Args:
        obj (Any): The object to check.

    Returns:
        bool: Whether the object is a DeviceView.
    """
    from polygraphy.cuda import DeviceView

    return isinstance(obj, DeviceView)


# The current design dispatches to the correct function implementation separately for each function call.
# Obviously, this has some performance cost and an alternative approach would be a more familiar inheritance
# pattern wherein we would have a BaseArray class and then child classes like NumpyArray, TorchArray, PolygraphyDeviceArray etc.
# That way, the dispatching logic would only have to run once when we construct an instance of one of these
# classes.
#
# The tradeoff is that the caller would then have to be careful that they are *not* passing in NumPy arrays,
# Torch tensors etc. directly, but have first wrapped them appropriately. Plus, at the interface boundaries,
# we would have to unwrap them once again since we don't want to expose the wrappers at the API level (the user
# should be able to work directly with NumPy arrays, PyTorch tensors etc.).
#
# To illustrate this a bit better, consider the two possible workflows:
#
# Option 1 (dispatch logic in each function, current design):
#
# def my_api_func(obj)
#     nbytes = util.array.nbytes(obj) # Dispatching logic needs to run on each function call
#     dtype = util.array.dtype(obj)
#     # Do something interesting, then...
#     return obj
#
# Option 2 (class hierarchy, possible alternative design):
#
# # Assume we have:
#
# class BaseArray:
#     ...
#
# class TorchArray:
#     ...
#
# # etc.
#
# def my_api_func()
#     obj = wrap_array(obj) # Dispatch logic only runs once
#     nbytes = obj.nbytes
#     dtype = obj.dtype
#     # Do something interesting, then...
#     return unwrap_array(obj) # Need to return the np.ndarray/torch.Tensor/DeviceView, *not* the wrapper
#
# In Polygraphy, the number of calls to `wrap_array`/`unwrap_array` would most likely be quite high
# relative to the number of calls to the actual methods, so the perfomance hit of the current implementation
# may not be that significant. If it is, then it should be straightforward, though time-consuming, to switch to Option 2.
#
def dispatch(num_arrays=1):
    """
    Decorator that will dispatch to functions specific to a framework type, like NumPy or PyTorch,
    based on the type of the input.

    The decorated function should return a dictionary with implementations for all supported types.
    The following keys may be specified: ["torch", "numpy", "device_view", "number"].

    Args:
        num_arrays (int):
            The number of arrays expected.
            The naming convention for the array arguments is as follows:
            - For a single array, the argument is called "obj".
            - For two arrays, the arguments are called "lhs" and "rhs".
            - For N>2 arrays, the arguments are called "obj0", "obj1", ... "obj<N-1>"
            In the case of more than one array, this function will automatically convert the rest to be of the
            same kind as the first.
    """

    def dispatch_impl(func):
        def _get_key(obj):
            key = None

            if is_device_view(obj):
                key = "device_view"
            elif is_numpy(obj):
                key = "numpy"
            elif is_torch(obj):
                key = "torch"
            elif isinstance(obj, numbers.Number):
                key = "number"

            if not key:
                G_LOGGER.critical(
                    f"Function: {func.__name__} is unsupported for objects of type: {type(obj).__name__}"
                )
            return key

        if num_arrays < 0:
            G_LOGGER.critical(
                f"Function: {func.__name__} is unsupported with {num_arrays} < 0"
            )

        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            if len(args) < num_arrays:
                G_LOGGER.critical(
                    f"Function: {func.__name__} is unsupported for less than {num_arrays} positional arguments"
                )

            mapping = func()
            obj0 = args[0]
            key = _get_key(obj0)

            if key not in mapping:
                G_LOGGER.critical(
                    f"Function: {func.__name__} is unsupported for objects of type: {type(obj0).__name__}"
                )

            # Note that we can use to_torch/to_numpy here without a circular dependency because those functions
            # take the num_arrays=1 path.
            def convert_array(obj):
                if key == "torch":
                    return to_torch(obj)
                elif key == "numpy":
                    return to_numpy(obj)
                else:
                    G_LOGGER.critical(
                        f"Function: {func.__name__} is unsupported for objects of type: {type(obj).__name__}"
                    )

            converted_args = (
                [obj0]
                + list(map(convert_array, args[1:num_arrays]))
                + list(args[num_arrays:])
            )

            return mapping[key](*converted_args, **kwargs)

        return wrapped

    return dispatch_impl


##
## Conversion Functions
##


@mod.export()
@dispatch()
def to_torch():
    """
    Converts an array or tensor to a PyTorch tensor.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        torch.Tensor: The PyTorch tensor.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj,
        "numpy": lambda obj: torch.from_numpy(obj),
        "number": lambda obj: torch.tensor(obj),
    }


@mod.export()
@dispatch()
def to_numpy():
    """
    Converts an array or tensor to a NumPy array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        np.ndarray: The NumPy array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.numpy(force=True),
        "numpy": lambda obj: obj,
        "number": lambda obj: np.array(obj),
    }


##
## Metadata
##


@mod.export()
@dispatch()
def nbytes():
    """
    Calculate the number of bytes required by the input array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        int: The number of bytes required by the array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.nelement() * obj.element_size(),
        "numpy": lambda obj: obj.nbytes,
        "device_view": lambda obj: obj.nbytes,
    }


@mod.export()
@dispatch()
def size():
    """
    Calculate the volume of the input array

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        int: The volume of the array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.numel(),
        "numpy": lambda obj: obj.size,
    }


@mod.export()
@dispatch()
def data_ptr():
    """
    Return a pointer to the first element of the input array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        int: A pointer to the first element of the array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.data_ptr(),
        "numpy": lambda obj: obj.ctypes.data,
        "device_view": lambda obj: obj.ptr,
    }


@mod.export()
@dispatch()
def is_on_cpu():
    """
    Returns whether the input array is in CPU memory.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        bool: Whether the array is in CPU, i.e. host, memory.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.device.type == "cpu",
        "numpy": lambda _: True,
        "device_view": lambda _: False,
    }


@mod.export()
@dispatch()
def is_on_gpu():
    """
    Returns whether the input array is in GPU memory.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        bool: Whether the array is in GPU, i.e. host, memory.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.device.type == "cuda",
        "numpy": lambda _: False,
        "device_view": lambda _: True,
    }


@mod.export()
@dispatch()
def dtype():
    """
    Return the data type the input array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        DataType: The data type of the array

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    func = lambda obj: DataType.from_dtype(obj.dtype)
    return {"torch": func, "numpy": func, "device_view": func}


@mod.export()
@dispatch()
def shape():
    """
    Return the shape the input array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray, DeviceView]: The shape of the array

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    func = lambda obj: obj.shape
    return {"torch": func, "numpy": func, "device_view": func}


@mod.export()
def view(obj, dtype, shape):
    """
    Return a view of the the input array with the given data type and shape.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]):
                The array or tensor. Must be contiguous.
        dtype (DataType): The data type to use for the view.
        shape (Sequence[int]): The shape to use for the view.

    Returns:
        Union[torch.Tensor, numpy.ndarray, DeviceView]: The view of the array

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    if not is_contiguous(obj):
        G_LOGGER.critical(f"Input array to view() must be contiguous in memory")

    if is_device_view(obj):
        return obj.view(shape=shape, dtype=dtype)

    dtype = (
        DataType.to_dtype(dtype, "numpy")
        if is_numpy(obj)
        else DataType.to_dtype(dtype, "torch")
    )
    return obj.reshape(-1).view(dtype).reshape(shape)


@mod.export()
@dispatch()
def is_contiguous():
    """
    Checks whether the provided array is contiguous in memory.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        bool: Whether the array is contiguous in memory.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "torch": lambda obj: obj.is_contiguous(),
        "numpy": lambda obj: obj.flags["C_CONTIGUOUS"],
        "device_view": lambda _: True,
    }


##
## Memory Management
##


@mod.export()
@dispatch()
def make_contiguous():
    """
    Makes an array contiguous if it's not already.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceView]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray, DeviceView]: The contiguous array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def impl_numpy(obj):
        if is_contiguous(obj):
            return obj
        return np.ascontiguousarray(obj)

    return {
        "torch": lambda obj: obj.contiguous(),
        "numpy": impl_numpy,
        "device_view": lambda obj: obj,
    }


@mod.export()
@dispatch()
def resize_or_reallocate():
    """
    Resizes the provided buffer, possibly reallocating the buffer.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray, DeviceArray]): The array or tensor.
        shape (Sequence[int]): The desired shape of the buffer.

    Returns:
        Union[torch.Tensor, numpy.ndarray, DeviceArray]: The resized buffer, possibly reallocated.
    """

    def numpy_impl(obj, shape):
        if shape != obj.shape:
            try:
                obj.resize(shape, refcheck=False)
            except ValueError as err:
                G_LOGGER.warning(
                    f"Could not resize NumPy array to shape: {shape}. "
                    f"Allocating a new array instead.\nNote: Error was: {err}"
                )
                obj = np.empty(shape, dtype=np.dtype(obj.dtype))
        return obj

    return {
        "numpy": numpy_impl,
        "torch": lambda obj, shape: obj.resize_(shape) if shape != obj.shape else obj,
        "device_view": lambda obj, shape: (
            obj.resize(shape) if shape != obj.shape else obj
        ),
    }


##
## Math Helpers
##


@mod.export()
@dispatch()
def cast():
    """
    Casts an array to the specified type.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        dtype (DataType): The type to cast to.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The casted array.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj, dtype: np.array(obj.astype(dtype.numpy())),
        "torch": lambda obj, dtype: obj.to(DataType.to_dtype(dtype, "torch")),
    }


@mod.export()
@dispatch()
def any():
    """
    Return whether any of the values in the provided array evaluate to True.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        bool: Whether any of the values in the array evaluate to True.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.any(obj),
        "torch": lambda obj: bool(torch.any(obj)),
    }


@mod.export()
@dispatch()
def all():
    """
    Return whether all of the values in the provided array evaluate to True.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        bool: Whether all of the values in the array evaluate to True.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.all(obj),
        "torch": lambda obj: bool(torch.all(obj)),
    }


@mod.export()
@dispatch(num_arrays=2)
def equal():
    """
    Returns whether two arrays are equal

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        bool: Whether the arrays are equal.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    return {
        "torch": lambda lhs, rhs: torch.equal(lhs, rhs),
        "numpy": lambda lhs, rhs: np.array_equal(lhs, rhs),
    }


@mod.export()
@dispatch(num_arrays=2)
def subtract():
    """
    Subtracts the second array from the first.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The difference.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    return {
        "torch": lambda lhs, rhs: lhs - rhs,
        "numpy": lambda lhs, rhs: np.array(lhs - rhs),
    }


@mod.export()
@dispatch(num_arrays=2)
def divide():
    """
    Divides the first array by the second.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The quotient.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    return {
        "torch": lambda lhs, rhs: lhs / rhs,
        "numpy": lambda lhs, rhs: lhs / rhs,
    }


@mod.export()
@dispatch(num_arrays=2)
def allclose():
    """
    Returns whether all the values in two arrays are within the given thresholds.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.
        rtol (float): The relative tolerance. Defaults to 1e-5.
        atol (float): The absolute tolerance. Defaults to 1e-8.

    Returns:
        bool: Whether the arrays are close.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    DEFAULT_RTOL = 1e-5
    DEFAULT_ATOL = 1e-8

    return {
        "torch": lambda lhs, rhs, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL: torch.allclose(
            lhs, rhs, rtol=rtol, atol=atol
        ),
        "numpy": lambda lhs, rhs, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL: np.allclose(
            lhs, rhs, rtol=rtol, atol=atol
        ),
    }


@mod.export()
def unravel_index(index, shape):
    """
    Unravels a flat index into a N-dimensional index based on the specified shape.

    Args:
        index (int): The flat index.
        shape (Sequence[int]): The shape on which to unravel the index.

    Returns:
        Tuple[int]: The N-dimensional index.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    index = int(index)

    nd_index = []
    for dim in reversed(shape):
        nd_index.insert(0, index % dim)
        index = index // dim

    return tuple(nd_index)


@mod.export()
@dispatch()
def histogram():
    """
    Compute a histogram for the given array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        range (Tuple[float, float]): The lower and upper range of the bins.

    Returns:
        Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]:
            The histogram values and the bin edges

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj, range=None):
        # PyTorch doesn't support histograms for all types, so cast to FP32
        original_dtype = obj.dtype
        hist, bins = torch.histogram(obj.to(torch.float32), bins=10, range=range)
        return hist.to(original_dtype), bins.to(original_dtype)

    return {
        "numpy": lambda obj, range=None: np.histogram(obj, bins=10, range=range),
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def max():
    """
    Returns the maximum value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The maximum value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.amax(obj).item(),
        "torch": lambda obj: torch.max(obj).item(),
    }


@mod.export()
@dispatch()
def argmax():
    """
    Returns the flattened index of the maximum value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        int: The flattened index.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj):
        # Torch argmax doesn't support bools
        return torch.argmax(obj.to(torch.float32))

    return {
        "numpy": lambda obj: np.argmax(obj),
        "torch": lambda obj: torch_impl(obj),
    }


@mod.export()
@dispatch()
def min():
    """
    Returns the minimum value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The minimum value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.amin(obj).item(),
        "torch": lambda obj: torch.min(obj).item(),
    }


@mod.export()
@dispatch()
def argmin():
    """
    Returns the flattened index of the minimum value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        int: The flattened index.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj):
        # Torch argmin doesn't support bools
        return torch.argmin(obj.to(torch.float32))

    return {
        "numpy": lambda obj: np.argmin(obj),
        "torch": lambda obj: torch_impl(obj),
    }


@mod.export()
@dispatch()
def mean():
    """
    Returns the mean value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        dtype (DataType): The mean compute type.

    Returns:
        Any: The mean value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj, dtype=None: np.mean(
            obj, dtype=DataType.to_dtype(dtype, "numpy") if dtype is not None else None
        ),
        "torch": lambda obj, dtype=None: torch.mean(
            obj, dtype=DataType.to_dtype(dtype, "torch") if dtype is not None else None
        ),
    }


@mod.export()
@dispatch()
def std():
    """
    Returns the standard deviation of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The standard deviation

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj):
        # torch.var is only supported for floats, so cast up and then back.
        obj_fp32 = obj.to(torch.float32)
        try:
            return torch.std(obj_fp32, correction=0)
        except AttributeError:
            return torch.std(obj_fp32, unbiased=False)

    return {
        "numpy": lambda obj: np.std(obj),
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def var():
    """
    Returns the variance of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The variance

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj):
        # torch.var is only supported for floats, so cast up and then back.
        obj_fp32 = obj.to(torch.float32)
        try:
            return torch.var(obj_fp32, correction=0)
        except AttributeError:
            return torch.var(obj_fp32, unbiased=False)

    return {
        "numpy": lambda obj: np.var(obj),
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def median():
    """
    Returns the median value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The median value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_impl(obj):
        # Median in PyTorch doesn't work as expected for arrays with an even number of elements - instead
        # of returning the average of the two middle elements, it just returns the smaller one.
        # It is also not implemented for some types, so cast to FP32 for compute.

        original_dtype = obj.dtype
        obj = obj.to(torch.float32)

        rv = 0
        if obj.nelement() % 2 == 1:
            rv = torch.median(obj)
        else:
            smaller = torch.median(obj)
            larger = torch.median(torch.cat([obj.flatten(), torch.max(obj)[None]]))
            rv = (smaller + larger) / 2.0
        return rv.to(original_dtype)

    return {
        "numpy": lambda obj: np.median(obj),
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def quantile():
    """
    Returns the value of the q quantile of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        q  (float): Quantile to compute, expected range [0, 1]

    Returns:
        Any: The quantile value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def numpy_impl(obj, q):
        if obj.size == 0:
            return np.inf
        return np.quantile(obj, q)

    def torch_impl(obj, q):
        if obj.numel() == 0:
            return torch.inf
        original_dtype = obj.dtype
        obj = obj.to(torch.float32)
        qunatile_val = torch.quantile(obj, q)
        return qunatile_val.to(original_dtype)

    return {
        "numpy": numpy_impl,
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def topk():
    """
    Returns a tuple of the top k values and indices of an array along a specified axis.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        k (int): The number of values to return. This is clamped to the length of obj along the given axis.
        axis (int): The axis to perform the topk computation on

    Returns:
        Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]: A tuple containing a pair of arrays,
            the first being the values and the second being the indices of the top k values along the specified axis

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def numpy_impl(obj, k, axis):
        # NumPy doesn't have a Top K implementation
        indices = np.argsort(-obj, axis=axis, kind="stable")
        axis_len = indices.shape[axis]
        indices = np.take(indices, np.arange(0, builtins.min(k, axis_len)), axis=axis)
        return np.take_along_axis(obj, indices, axis=axis), indices

    def torch_impl(obj, k, axis):
        axis_len = obj.shape[axis]

        # Top K has no implementation for float16 in torch-cpu, so
        # If gpu is available, run computation there
        # Otherwise, run the calculation on cpu using fp32 precision
        if obj.dtype == torch.float16:
            if torch.cuda.is_available():
                original_device = obj.device
                ret = tuple(
                    torch.topk(obj.to("cuda"), builtins.min(k, axis_len), dim=axis)
                )
                return (ret[0].to(original_device), ret[1].to(original_device))
            else:
                ret = tuple(
                    torch.topk(
                        obj.type(torch.float32), builtins.min(k, axis_len), dim=axis
                    )
                )
                return (ret[0].type(torch.float16), ret[1].type(torch.float16))
        return tuple(torch.topk(obj, builtins.min(k, axis_len), dim=axis))

    return {
        "numpy": numpy_impl,
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def abs():
    """
    Returns the absolute value of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Any: The absolute value

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """

    def torch_abs_impl(obj):
        # PyTorch doesn't support abs for all types, so cast to FP32
        original_dtype = obj.dtype
        return torch.abs(obj.to(torch.float32)).to(original_dtype)

    return {
        "numpy": lambda obj: np.array(np.abs(obj)),
        "torch": lambda obj: torch_abs_impl(obj),
    }


@mod.export()
@dispatch()
def isfinite():
    """
    Returns a boolean array indicating if each element of obj is finite or not.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The boolean array indicating which elements of obj are finite.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.isfinite(obj),
        "torch": lambda obj: torch.isfinite(obj),
    }


@mod.export()
@dispatch()
def isinf():
    """
    Returns a boolean array indicating if each element of obj is infinite or not.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The boolean array indicating which elements of obj are infinite.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.isinf(obj),
        "torch": lambda obj: torch.isinf(obj),
    }


@mod.export()
@dispatch()
def isnan():
    """
    Returns a boolean array indicating if each element of obj is NaN or not.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The boolean array indicating which elements of obj are NaN.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.isnan(obj),
        "torch": lambda obj: torch.isnan(obj),
        "number": lambda obj: math.isnan(obj),
    }


@mod.export()
@dispatch()
def argwhere():
    """
    Returns a indices of non-zero array elements

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: An (N, obj.ndim) array containing indices of non-zero elements of obj

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.argwhere(obj),
        "torch": lambda obj: torch.argwhere(obj),
    }


@mod.export()
@dispatch()
def ravel():
    """
    Flattens the input array

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The flattened input tensor

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.ravel(obj),
        "torch": lambda obj: torch.ravel(obj),
    }


@mod.export()
@dispatch()
def logical_not():
    """
    Computes the logical not of an array

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]):
                The input array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The logical not.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.logical_not(obj),
        "torch": lambda obj: torch.logical_not(obj),
    }


@mod.export()
@dispatch(num_arrays=2)
def logical_xor():
    """
    Computes the logical exclusive-or of two arrays.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The logical xor.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda lhs, rhs: np.logical_xor(lhs, rhs),
        "torch": lambda lhs, rhs: torch.logical_xor(lhs, rhs),
    }


@mod.export()
@dispatch(num_arrays=2)
def logical_and():
    """
    Computes the logical and of two arrays.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The logical and.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda lhs, rhs: np.logical_and(lhs, rhs),
        "torch": lambda lhs, rhs: torch.logical_and(lhs, rhs),
    }


@mod.export()
@dispatch(num_arrays=2)
def greater():
    """
    Returns a boolean array indicating where lhs is greater than rhs

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: Boolean array indicating whether lhs > rhs.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda lhs, rhs: np.greater(lhs, rhs),
        "torch": lambda lhs, rhs: torch.gt(lhs, rhs),
    }


@mod.export()
@dispatch(num_arrays=3)
def where():
    """
    Returns an array containing elements from lhs when cond is true, and rhs when cond is false.
    Computes the logical and of two arrays.

    Args:
        cond (Union[torch.Tensor, numpy.ndarray]):
                The condition array or tensor.
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: Selected elements from lhs if cond is true, and rhs otherwise

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda cond, lhs, rhs: np.where(cond, lhs, rhs),
        "torch": lambda cond, lhs, rhs: torch.where(cond, lhs, rhs),
    }


@mod.export()
@dispatch(num_arrays=2)
def power():
    """
    Computes the element-wise power of an array to the given exponent.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]):
                The base array or tensor.
        exponent (Union[int, float, torch.Tensor, numpy.ndarray]):
                The exponent value or array.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The power result.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj, exponent: np.power(obj, exponent),
        "torch": lambda obj, exponent: torch.pow(obj, exponent),
    }


@mod.export()
@dispatch()
def sum():
    """
    Computes the sum of all elements in the array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[Number, torch.Tensor, numpy.ndarray]: The sum of all elements.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.sum(obj),
        "torch": lambda obj: torch.sum(obj),
    }


@mod.export()
@dispatch()
def sqrt():
    """
    Computes the element-wise square root of an array.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The square root results.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj: np.sqrt(obj),
        "torch": lambda obj: torch.sqrt(obj),
    }


@mod.export()
@dispatch(num_arrays=2)
def multiply():
    """
    Computes the element-wise multiplication of two arrays.

    Args:
        lhs (Union[torch.Tensor, numpy.ndarray]):
                The first array or tensor.
        rhs (Union[torch.Tensor, numpy.ndarray]):
                The second array or tensor.

    Returns:
        Union[torch.Tensor, numpy.ndarray]: The element-wise product.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda lhs, rhs: np.multiply(lhs, rhs),
        "torch": lambda lhs, rhs: torch.mul(lhs, rhs),
    }


@mod.export()
@dispatch()
def pad():
    """
    Pads the input array on specified axes with number of pre/post padding values.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        pads (List[Tuple(int)]): The padding values for each axis
        axes (List(int)): Optional: the axes that the pads argument corresponds to. Defaults to None.
        constant_val (sequence or scalar): Optional: The value to pad for "constant" mode padding. Defaults to 0.
        mode (str): The padding mode (Default - "constant")
    """

    def numpy_impl(obj, pads, axes=None, mode="constant", constant_val=0):
        if axes is None:
            if len(pads) != len(obj.shape):
                raise ValueError(
                    f"Number of padding tuples ({len(pads)}) must match number of dimensions ({len(obj.shape)})"
                )
            return np.pad(obj, pads, mode, constant_values=constant_val)
        else:
            if len(axes) != len(pads):
                raise ValueError("Number of padding tuples and axes must be the same!")
            final_pads = [(0, 0) for _ in range(len(obj.shape))]
            for axis, pad_vals in zip(axes, pads):
                final_pads[axis] = pad_vals
            return np.pad(obj, final_pads, mode, constant_values=constant_val)

    # torch.nn.Functional.pad starts with the last axis first, and expects a singular list of all pad values
    # e.g. (last_axis_prepad, last_axis_postpad, last_axis-1_prepad ...)
    def torch_impl(obj, pads, axes=None, mode="constant", constant_val=0):
        if axes is None:
            if len(pads) != len(obj.shape):
                raise ValueError(
                    f"Number of padding tuples ({len(pads)}) must match number of dimensions ({len(obj.shape)})"
                )
            # Convert list of tuple of padding values into list of padding values, and reverse it.
            torch_pads = [x for pad_tups in pads for x in pad_tups]
            torch_pads.reverse()
            return torch.nn.functional.pad(obj, torch_pads, mode, constant_val)
        else:
            if len(axes) != len(pads):
                raise ValueError("Number of padding tuples and axes must be the same!")
            # Create default padding list
            final_pads = [(0, 0) for _ in range(len(obj.shape))]
            for axis, pad_vals in zip(axes, pads):
                final_pads[axis] = pad_vals
            # Similarly convert this into a list and reverse it for torch.
            final_pads = [x for pad_tups in final_pads for x in pad_tups]
            final_pads.reverse()
            return torch.nn.functional.pad(obj, final_pads, mode, constant_val)

    return {
        "numpy": numpy_impl,
        "torch": torch_impl,
    }


@mod.export()
@dispatch()
def reshape():
    """
    Reshapes the input tensor to the provided shape

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        shape (Union[tuple(ints), List(ints)]): The new shape of the tensor.

    Returns:
        Union[Number, torch.Tensor, numpy.ndarray]: The reshaped tensor.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj, shape: np.reshape(obj, shape),
        "torch": lambda obj, shape: torch.reshape(obj, shape),
    }


@mod.export()
@dispatch()
def transpose():
    """
    Transposes the array according to the provided axes.

    Args:
        obj (Union[torch.Tensor, numpy.ndarray]): The array or tensor.
        axes (Union[tuple(ints), List(ints)]): The axes to transpose.

    Returns:
        Union[Number, torch.Tensor, numpy.ndarray]: The transposed tensor.

    Raises:
        PolygraphyException: if the input is of an unrecognized type.
    """
    return {
        "numpy": lambda obj, axes: np.transpose(obj, axes),
        "torch": lambda obj, axes: torch.permute(obj, axes),
    }