# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""BufferView and Python Array Protocol interop."""

from typing import Literal, Optional, Tuple, cast
import logging
import numpy as np
import numpy.lib.mixins

from ._binding import (
    BufferCompatibility,
    BufferUsage,
    HalBufferView,
    HalDevice,
    HalElementType,
    MappedMemory,
    MemoryType,
    HalFence,
)
from .dtypes import map_dtype_to_hal_element_type

__all__ = [
    "asdevicearray",
    "DeviceArray",
]

_DEVICE_HANDLED_FUNCTIONS = {}


def _device_implements(np_function):
    """Decorator that registers a base class implementation."""

    def decorator(func):
        _DEVICE_HANDLED_FUNCTIONS[np_function] = func
        return func

    return decorator


class DeviceArray(numpy.lib.mixins.NDArrayOperatorsMixin):
    """An IREE device array.

    Device arrays can be in one of two states:

    1. Host accessible: The array will be backed by host accessible memory
       and can have the usual things done with it that one expects to be
       able to do with an ndarray.
    2. Device resident: The array is just a handle to a device resident
       Buffer (and BufferView wrapper). Metadata about the array are accessible
       (shape and dtype) but anything that touches the data cannot be accessed
       in this state.

    How a device array comes into existence controls how it can transition
    between these states:

    * A user can create a DeviceArray explicitly with a device allocator.
      Such an array will not be implicitly convertible to host accessible,
      although accessors exist to do so.
    * When created by the platform with a synchronization policy, then
      implicit transfer back to the host will trigger appropriate waits and
      be performed automatically (this is the common case for function return
      values if not otherwise configured, as an example).
    """

    def __init__(
        self,
        device: HalDevice,
        buffer_view: HalBufferView,
        implicit_host_transfer: bool = False,
        override_dtype=None,
    ):
        self._device = device
        self._buffer_view = buffer_view
        self._implicit_host_transfer = implicit_host_transfer
        self._override_dtype = override_dtype

        # If the array is host accessible, these will be non-None.
        self._mapped_memory: Optional[MappedMemory] = None
        self._host_array: Optional[np.ndarray] = None

    def __array__(self, dtype=None):
        host_array = self._transfer_to_host(True)
        if dtype is None:
            return host_array
        else:
            return host_array.__array__(dtype)  # pytype: disable=attribute-error

    def __array_function__(self, func, types, args, kwargs):
        if func in _DEVICE_HANDLED_FUNCTIONS:
            return _DEVICE_HANDLED_FUNCTIONS[func](*args, **kwargs)

        # Anything else forces a transfer to host and then delegates to the
        # host array.
        host_array = self.to_host()
        return host_array.__array_function__(
            func, types, args, kwargs
        )  # pytype: disable=attribute-error

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        def transfer_input(inp):
            if isinstance(inp, DeviceArray):
                return inp._transfer_to_host(True)
            return inp

        host_inputs = [transfer_input(inp) for inp in inputs]
        if kwargs.get("out") is not None:
            kwargs = dict(kwargs)
            kwargs["out"] = tuple(transfer_input(inp) for inp in kwargs["out"])
        return getattr(ufunc, method)(*host_inputs, **kwargs)

    def __repr__(self):
        return f"<IREE DeviceArray: shape={np.shape(self)}, dtype={self.dtype}>"

    @property
    def is_host_accessible(self):
        """Whether this array is currently host accessible."""
        return self._host_array is not None

    def to_host(self) -> np.ndarray:
        """Return the array as host accessible NumPy ndarray.
        This may map the memory or create a copy depending on whether the array is
        mappable to the host."""
        return self._transfer_to_host(False)

    def _is_mappable(self) -> bool:
        buffer = self._buffer_view.get_buffer()
        if buffer.memory_type() & int(MemoryType.HOST_VISIBLE) != int(
            MemoryType.HOST_VISIBLE
        ):
            return False
        if buffer.allowed_usage() & int(BufferUsage.MAPPING_SCOPED) != int(
            BufferUsage.MAPPING_SCOPED
        ):
            return False
        return True

    def _transfer_to_host(self, implicit) -> np.ndarray:
        if implicit and not self._implicit_host_transfer:
            raise ValueError(
                "DeviceArray cannot be implicitly transferred to the host: "
                "if necessary, do an explicit transfer via .to_host()"
            )
        if self._is_mappable():
            if self._host_array is None:
                self._mapped_memory, self._host_array = self._map_to_host()
            return self._host_array
        else:
            return self._copy_to_host()

    def _map_to_host(self) -> Tuple[MappedMemory, np.ndarray]:
        # TODO: When synchronization is enabled, need to block here.
        raw_dtype = self._get_raw_dtype()
        mapped_memory = self._buffer_view.map()
        host_array = cast(
            np.ndarray, mapped_memory.asarray(self._buffer_view.shape, raw_dtype)
        )
        # Detect if we need to force an explicit conversion. This happens when
        # we were requested to pretend that the array is in a specific dtype,
        # even if that is not representable on the device. You guessed it:
        # this is to support bools.
        if self._override_dtype is not None and self._override_dtype != raw_dtype:
            host_array = host_array.astype(self._override_dtype)
        return mapped_memory, host_array

    def _copy_to_host(self) -> np.ndarray:
        # TODO: When synchronization is enabled, need to block here.
        raw_dtype = self._get_raw_dtype()
        source = self._buffer_view.get_buffer()
        alloc = self._device.allocator
        mem_type = int(MemoryType.HOST_LOCAL | MemoryType.DEVICE_VISIBLE)
        import_usage = int(BufferUsage.TRANSFER_TARGET | BufferUsage.MAPPING_PERSISTENT)

        def dma_and_wait(host_buf):
            sem = self._device.create_semaphore(0)
            self._device.queue_copy(
                source,
                host_buf,
                wait_semaphores=HalFence.create_at(sem, 0),
                signal_semaphores=HalFence.create_at(sem, 1),
            )
            HalFence.create_at(sem, 1).wait()

        # Check if the allocator supports importing host memory (HIP/CUDA do,
        # Vulkan's native allocator does not). On HIP the import path avoids
        # a staging allocation and is 1.3-13x faster depending on size.
        compat = alloc.query_buffer_compatibility(
            memory_type=mem_type,
            allowed_usage=import_usage,
            intended_usage=import_usage,
            allocation_size=source.byte_length(),
        )
        if compat & int(BufferCompatibility.IMPORTABLE):
            result = np.empty(self._buffer_view.shape, dtype=raw_dtype)
            host_buf = alloc.import_host_allocation(
                result,
                memory_type=mem_type,
                allowed_usage=import_usage,
            )
            dma_and_wait(host_buf)
        else:
            staging_usage = int(
                BufferUsage.TRANSFER_TARGET | BufferUsage.MAPPING_SCOPED
            )
            host_buf = alloc.allocate_buffer(
                memory_type=mem_type,
                allowed_usage=staging_usage,
                allocation_size=source.byte_length(),
            )
            dma_and_wait(host_buf)
            mapped = host_buf.map()
            result = mapped.asarray(self._buffer_view.shape, raw_dtype)

        if self._override_dtype is not None and self._override_dtype != raw_dtype:
            return result.astype(self._override_dtype)
        return result

    def _get_raw_dtype(self):
        return HalElementType.map_to_dtype(self._buffer_view.element_type)

    @property
    def dtype(self):
        if self._override_dtype:
            return self._override_dtype
        return self._get_raw_dtype()

    @property
    def shape(self):
        return np.shape(self)

    def astype(
        self,
        dtype,
        casting: Literal["no", "equiv", "safe", "same_kind", "unsafe"] = "unsafe",
        copy=True,
    ):
        if self.dtype == dtype and not copy:
            return self
        host_ary = self.to_host()
        return host_ary.astype(dtype, casting=casting, copy=copy)

    def reshape(self, *args):
        # TODO(scotttodd): add a native impl with a new buffer_view of the same data
        # TODO(scotttodd): return DeviceArray instead of host ndarray?
        host_ary = self.to_host()
        return host_ary.reshape(*args)

    def __iter__(self):
        host_ary = self.to_host()
        return host_ary.__iter__()

    def __getitem__(self, index):
        host_ary = self.to_host()
        return host_ary.__getitem__(index)

    def __deepcopy__(self, memo):
        return self.to_host()

    def __reduce__(self):
        # We just reduce to the actual host ndarray, which supports the necessary
        # serialization protocols.
        return _restore_reduced_array, (self.to_host(),)


def _restore_reduced_array(ary):
    return ary


# Function implementations with custom behavior.
@_device_implements(np.shape)
def _(arr: DeviceArray):
    return arr._buffer_view.shape


@_device_implements(np.reshape)
def _(arr: DeviceArray, *args):
    return arr.reshape(*args)


def asdevicearray(
    device: HalDevice,
    a,
    dtype=None,
    *,
    implicit_host_transfer: bool = False,
    memory_type=MemoryType.DEVICE_LOCAL,
    allowed_usage=BufferUsage.DEFAULT,
    element_type: Optional[HalElementType] = None,
) -> DeviceArray:
    """Helper to create a DeviceArray from an arbitrary array like.

    This is similar in purpose and usage to np.asarray, except that it takes
    a device as the first argument. This may not be the best mechanism for
    getting a DeviceArray, depending on your use case, but it is reliable
    and simple. This function may make a defensive copy or cause implicit
    transfers to satisfy the request. The default creates a device-local buffer;
    callers that require direct host mapping should request a host-visible
    memory type and mapping usage explicitly. If this is important to you, then
    a lower level API is likely more appropriate.

    Note that additional flags `memory_type`, `allowed_usage` and `element_type`
    are only hints if creating a new DeviceArray. If `a` is already a DeviceArray,
    they are ignored.
    """
    if isinstance(a, DeviceArray):
        if dtype is None:
            return a
        # Need to do a conversion, which we currently do not support on the
        # device, so transfer back to the host.
        logging.warn(
            "Implicit dtype conversion of a DeviceArray forces a host transfer"
        )
    # First get an ndarray. Needs to be C-contiguous, enforcing it here.
    a = np.asarray(a, dtype=dtype, order="C")
    element_type = map_dtype_to_element_type(a.dtype)
    if element_type is None:
        raise ValueError(f"Could not map dtype {a.dtype} to IREE element type")
    buffer_view = device.allocator.allocate_buffer_copy(
        memory_type=memory_type,
        allowed_usage=allowed_usage,
        device=device,
        buffer=a,
        element_type=element_type,
    )
    if not isinstance(buffer_view, HalBufferView):
        raise TypeError("Expected allocate_buffer_copy to return a HalBufferView")
    return DeviceArray(
        device,
        buffer_view,
        implicit_host_transfer=implicit_host_transfer,
        override_dtype=a.dtype,
    )


def map_dtype_to_element_type(dtype) -> Optional[HalElementType]:
    return map_dtype_to_hal_element_type(dtype)
