"""Abstraction layer over PyOpenCL to present an interface that is common
between CUDA and OpenCL. This is only the subset that is needed so far. New
functionality can be added, but should be kept in sync with
:mod:`katsdpsigproc.cuda`.
"""

import pyopencl
import pyopencl.array
import numpy as np


class _DummyArray(object):
    """Trivial dummy for :class:`pyopencl.array.Array`, that just has
    a `data` attribute.
    """
    def __init__(self, data):
        self.data = data


class _PinnedAMD(np.ndarray):
    """Pinned memory abstraction for AMD GPUs. Copies to the device
    are done by unmapping the buffer, enqueuing the copy, and remapping
    the buffer. This is based on AMD's optimization guide.
    """
    def __new__(cls, context, queue, shape, dtype):
        dtype = np.dtype(dtype)
        n_bytes = int(np.product(shape)) * dtype.itemsize
        # Do not add READ or WRITE to the flags: doing so seems to cause AMD
        # drivers to allocate GPU memory.
        buffer = pyopencl.Buffer(
            context,
            pyopencl.mem_flags.ALLOC_HOST_PTR,
            n_bytes)
        array = pyopencl.enqueue_map_buffer(
            queue,
            buffer,
            pyopencl.map_flags.READ | pyopencl.map_flags.WRITE,
            0, shape, dtype)[0]
        mapping = array.base
        array = array.view(_PinnedAMD)
        array._buffer = buffer
        array._mapping = mapping
        array._array = _DummyArray(buffer)
        return array

    def _enqueue_transfer(self, queue, blocking, transfer):
        if not hasattr(self, '_buffer'):
            raise ValueError('cannot copy to/from a _PinnedAMD view')
        self._mapping.release(queue._pyopencl_command_queue)
        transfer()
        # The shape and dtype don't matter here as long as we map the entire buffer,
        # because we're going to throw away the array wrapper and just keep the
        # MemoryMap object.
        array = pyopencl.enqueue_map_buffer(
            queue._pyopencl_command_queue,
            self._buffer, pyopencl.map_flags.READ | pyopencl.map_flags.WRITE,
            0, (self._buffer.size,), np.uint8, is_blocking=blocking)[0]
        self._mapping = array.base

    def enqueue_write_buffer(self, queue, buffer, blocking):
        """Enqueue a copy from the host buffer to device memory. The host
        memory must not be touched while the copy is in progress.

        Parameters
        ----------
        queue : :class:`CommandQueue`
            Queue for the copy
        buffer : `pyopencl.array.Array`
            Target buffer
        blocking : bool
            If `True`, wait for the copy to complete before returning
        """
        def transfer():
            queue.enqueue_copy_buffer(self._array, buffer)
        self._enqueue_transfer(queue, blocking, transfer)

    def enqueue_read_buffer(self, queue, buffer, blocking):
        """Enqueue a copy to the host buffer from device memory. The host
        memory must not be touched while the copy is in progress.

        Parameters
        ----------
        queue : :class:`CommandQueue`
            Queue for the copy
        buffer : `pyopencl.array.Array`
            Source buffer
        blocking : bool
            If `True`, wait for the copy to complete before returning
        """
        def transfer():
            queue.enqueue_copy_buffer(buffer, self._array)
        self._enqueue_transfer(queue, blocking, transfer)

    def enqueue_write_buffer_rect(
            self, queue, buffer, buffer_origin, data_origin, shape,
            buffer_strides, data_strides, blocking=True):
        def transfer():
            queue.enqueue_copy_buffer_rect(
                self._array, buffer, data_origin, buffer_origin, shape,
                data_strides, buffer_strides)
        self._enqueue_transfer(queue, blocking, transfer)

    def enqueue_read_buffer_rect(
            self, queue, buffer, buffer_origin, data_origin, shape,
            buffer_strides, data_strides, blocking=True):
        def transfer():
            queue.enqueue_copy_buffer_rect(
                buffer, self._array, buffer_origin, data_origin, shape,
                buffer_strides, data_strides)
        self._enqueue_transfer(queue, blocking, transfer)


class Program(object):
    """Abstraction of a program object"""
    def __init__(self, pyopencl_program):
        self._pyopencl_program = pyopencl_program

    def get_kernel(self, name):
        """Create a new kernel

        Parameters
        ----------
        name : str
            Name of the kernel function
        """
        return Kernel(self, name)


class Kernel(object):
    """Abstraction of a kernel object. The object can be enqueued using
    :meth:`CommandQueue.enqueue_kernel`.

    The recommended way to create this object is via
    :meth:`Program.get_kernel`.
    """
    def __init__(self, program, name):
        self._pyopencl_kernel = pyopencl.Kernel(program._pyopencl_program, name)


class Event(object):
    """Abstraction of an event. This is more akin to a CUDA event than an
    OpenCL event, in that it is a marker in a command queue rather than
    associated with a specific command.
    """
    def __init__(self, pyopencl_event):
        self._pyopencl_event = pyopencl_event

    def wait(self):
        """Block until the event has completed"""
        self._pyopencl_event.wait()

    def time_since(self, prior_event):
        """Return the time in seconds from `prior_event` to self. Unlike the
        PyCUDA method of the same name, this will wait for the events to
        complete if they have not already.
        """
        prior_event.wait()
        self.wait()
        return 1e-9 * (self._pyopencl_event.profile.start
                       - prior_event._pyopencl_event.profile.end)

    def time_till(self, next_event):
        """Return the time in seconds from this event to `next_event`. See
        :meth:`time_since`.
        """
        return next_event.time_since(self)


class Device(object):
    """Abstraction of an OpenCL device"""
    def __init__(self, pyopencl_device):
        self._pyopencl_device = pyopencl_device

    def make_context(self):
        """Create a new context associated with this device"""
        return Context(pyopencl.Context([self._pyopencl_device]))

    @property
    def name(self):
        """Return human-readable name for the device"""
        return self._pyopencl_device.name

    @property
    def platform_name(self):
        """Return human-readable name for the platform owning the device"""
        return self._pyopencl_device.platform.name

    @property
    def driver_version(self):
        """Return human-readable name for the driver version"""
        return self._pyopencl_device.driver_version

    @property
    def is_cuda(self):
        """Whether the device is a CUDA device"""
        return False

    @property
    def is_gpu(self):
        """Whether device is a GPU"""
        return self._pyopencl_device.type & pyopencl.device_type.GPU

    @property
    def is_accelerator(self):
        """Whether device is an accelerator (as defined by OpenCL device
        types)"""
        return self._pyopencl_device.type & pyopencl.device_type.ACCELERATOR

    @property
    def is_cpu(self):
        """Whether the device is a CPU"""
        return self._pyopencl_device.type & pyopencl.device_type.CPU

    @property
    def simd_group_size(self):
        """The number of workitems that run in lock-step. This must only
        be used to tune performance parameters; there are no guarantees about
        memory coherency, forward progress etc.
        """
        return pyopencl.characterize.get_simd_group_size(self._pyopencl_device, 4)

    @classmethod
    def _get_platforms(cls):
        """Return all platforms"""
        try:
            return pyopencl.get_platforms()
        except pyopencl.LogicError:
            # OpenCL considers it an error if there are no platforms
            # available.
            return []

    @classmethod
    def get_devices(cls):
        """Return a list of all devices on all platforms"""
        ans = []
        for platform in cls._get_platforms():
            for device in platform.get_devices():
                ans.append(Device(device))
        return ans

    @classmethod
    def get_devices_by_platform(cls):
        """Return a list of all devices, with a sub-list per platform."""
        ans = []
        for platform in cls._get_platforms():
            ans.append([Device(device) for device in platform.get_devices()])
        return ans


class Context(object):
    """Abstraction of an OpenCL context"""
    def __init__(self, pyopencl_context):
        self._pyopencl_context = pyopencl_context
        device = pyopencl_context.devices[0]
        self._internal_queue = pyopencl.CommandQueue(pyopencl_context, device)
        self._force_pinned_amd = False   # Used for unit tests

    @property
    def device(self):
        """Return the device associated with the context (or the first device,
        if there are multiple)."""
        return Device(self._pyopencl_context.devices[0])

    def compile(self, source, extra_flags=None):
        """Build a program object from source

        Parameters
        ----------
        source : str
            OpenCL source code
        extra_flags : list
            OpenCL-specific parameters to pass to the compiler
        """
        # source is passed through str because it might arrive as Unicode,
        # triggering a warning
        program = pyopencl.Program(self._pyopencl_context, str(source))
        program.build(extra_flags)
        return Program(program)

    def allocate_raw(self, n_bytes):
        """Create an untyped buffer on the device."""
        return pyopencl.Buffer(self._pyopencl_context, pyopencl.mem_flags.READ_WRITE, n_bytes)

    def allocate(self, shape, dtype, raw=None):
        """Create a typed buffer on the device.

        Parameters
        ----------
        shape : tuple
            Shape for the array
        dtype : numpy dtype
            Type for the data
        raw : PyOpenCL buffer, optional
            Memory backing the array (automatically allocated if `None`
        """
        return pyopencl.array.Array(self._pyopencl_context, shape, dtype, data=raw)

    def allocate_pinned(self, shape, dtype):
        """Create a buffer in host memory that can be efficiently copied
        to and from the device.

        Parameters
        ----------
        shape : tuple
            Shape for the array
        dtype : numpy dtype
            Type for the data
        """
        device = self._pyopencl_context.devices[0]
        if (self._force_pinned_amd
            or ((device.type & pyopencl.device_type.GPU)
                and device.platform.name == 'AMD Accelerated Parallel Processing')):
            return _PinnedAMD(self._pyopencl_context, self._internal_queue, shape, dtype)
        elif device.platform.name == 'NVIDIA CUDA':
            # Based on NVIDIA's recommendation: create a device buffer and
            # leave it mapped permanently.
            n_bytes = int(np.product(shape)) * dtype.itemsize
            buf = pyopencl.Buffer(
                self._pyopencl_context,
                pyopencl.mem_flags.ALLOC_HOST_PTR | pyopencl.mem_flags.READ_ONLY,
                n_bytes)
            return pyopencl.enqueue_map_buffer(
                self._internal_queue,
                buf,
                pyopencl.map_flags.READ | pyopencl.map_flags.WRITE,
                0, shape, dtype)[0]
        else:
            return np.empty(shape, dtype)

    def allocate_svm_raw(self, n_bytes):
        raise NotImplementedError("PyOpenCL does not support OpenCL Shared Virtual Memory")

    def allocate_svm(self, shape, dtype, raw=None):
        raise NotImplementedError("PyOpenCL does not support OpenCL Shared Virtual Memory")

    def create_command_queue(self, profile=False):
        """Create a new command queue associated with this context

        Parameters
        ----------
        profile : boolean
            If true, the command queue will support timing kernels
        """
        return CommandQueue(self, profile=profile)

    def create_tuning_command_queue(self):
        """Create a new command queue for doing autotuning"""
        return TuningCommandQueue(self)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False


class CommandQueue(object):
    """Abstraction of a command queue. If no existing command queue is passed
    to the constructor, a new one is created.

    Parameters
    ----------
    context : :class:`Context`
        Context owning the queue
    pyopencl_command_queue : :class:`pyopencl.CommandQueue` (optional)
        Existing command queue to wrap
    profile : boolean
        If true and `pyopencl_command_queue` is omitted, enabling profiling
        (timing) on the queue.
    """
    def __init__(self, context, pyopencl_command_queue=None, profile=False):
        self.context = context
        if pyopencl_command_queue is None:
            pyopencl_device = context._pyopencl_context.devices[0]
            properties = 0
            if profile:
                properties |= pyopencl.command_queue_properties.PROFILING_ENABLE
            pyopencl_command_queue = pyopencl.CommandQueue(
                context._pyopencl_context, pyopencl_device,
                properties)
        self._pyopencl_command_queue = pyopencl_command_queue

    def enqueue_read_buffer(self, buffer, data, blocking=True):
        """Copy data from the device to the host. Only whole-buffer copies are
        supported, and the shape and type must match. In general, one should use
        the convenience functions in :class:`accel.DeviceArray`.

        Parameters
        ----------
        buffer : `pyopencl.array.Array`
            Source
        data : array-like
            Target
        blocking : boolean (optional)
            If true (default) the call blocks until the copy is complete
        """
        if hasattr(data, 'enqueue_read_buffer'):
            data.enqueue_read_buffer(self, buffer, blocking)
        else:
            buffer.get(ary=data, queue=self._pyopencl_command_queue, async_=not blocking)

    def enqueue_write_buffer(self, buffer, data, blocking=True):
        """Copy data from the host to the device. Only whole-buffer copies are
        supported, and the shape and type must match. In general, one should
        use the convenience functions in :class:`accel.DeviceArray`.

        Parameters
        ----------
        buffer : :class:`pyopencl.array.Array`
            Target
        data : array-like
            Source
        blocking : boolean (optional)
            If true (default), the call blocks until the source has been fully
            read (it has no necessarily reached the device).
        """
        if hasattr(data, 'enqueue_write_buffer'):
            data.enqueue_write_buffer(self, buffer, blocking)
        else:
            buffer.set(ary=data, queue=self._pyopencl_command_queue, async_=not blocking)

    def enqueue_copy_buffer(
            self, src_buffer, dest_buffer):
        """Copy one buffer to another.

        Parameters
        ----------
        src_buffer,dest_buffer : :class:`pyopencl.array.Array`
            Source and destination buffers
        """
        pyopencl.enqueue_copy(
            self._pyopencl_command_queue,
            src=src_buffer.data, dest=dest_buffer.data)

    def enqueue_copy_buffer_rect(
            self, src_buffer, dest_buffer, src_origin, dest_origin,
            shape, src_strides, dest_strides):
        """Copy a subregion of one buffer to another. This is a low-level
        interface that ignores the shape, strides etc of the buffers, and
        treats them as byte arrays. It also only supports 3 or fewer
        dimensions. Use
        :py:meth:`~katsdpsigproc.accel.DeviceArray.copy_region` for a
        high-level interface.

        Parameters
        ----------
        src_buffer,dest_buffer : :class:`pyopencl.array.Array`
            Source and destination buffers
        src_origin,dest_origin : int
            Offsets for the start of the copy, in bytes
        shape : sequence of int
            Shape of the region to copy (1-3 elements). The first dimension is
            a byte count.
        src_strides,dest_strides : sequence of int
            Strides for the source and destination memory layout, with the same
            length as `shape`. The first element of each must be 1, and each
            element must be a factor of the next element.
        """
        assert src_strides[0] == 1
        assert dest_strides[0] == 1
        pyopencl.enqueue_copy(
            self._pyopencl_command_queue,
            src=src_buffer.data, dest=dest_buffer.data,
            src_origin=(src_origin,), dst_origin=(dest_origin,),
            region=shape,
            src_pitches=src_strides[1:], dst_pitches=dest_strides[1:])

    def enqueue_read_buffer_rect(
            self, buffer, data, buffer_origin, data_origin, shape,
            buffer_strides, data_strides, blocking=True):
        """Copy a region of a buffer to host memory. This is a low-level
        interface that ignores the shape, strides etc of the buffers, and
        treats them as byte arrays. It also only supports 3 or fewer dimensions. Use
        :py:meth:`~katsdpsigproc.accel.DeviceArray.set_region` for a high-level
        interface.

        Parameters
        ----------
        buffer : :class`pyopencl.array.Array`
            Source
        data : array-like
            Target
        buffer_origin, data_origin : int
            Offsets for the start of the copy, in bytes
        shape : sequence of int
            Shape of the region to copy (1-3 elements). The first dimension is
            a byte count.
        buffer_strides,data_strides : sequence of int
            Strides for the destination and source memory layout, with the same
            length as `shape`. The first element of each must be 1, and each
            element must be a factor of the next element.
        blocking : bool, optional
            If true, block until the transfer is complete.
        """
        assert buffer_strides[0] == 1
        assert data_strides[0] == 1
        if hasattr(data, 'enqueue_read_buffer_rect'):
            data.enqueue_read_buffer_rect(self, buffer, buffer_origin, data_origin, shape,
                                          buffer_strides, data_strides, blocking)
        else:
            pyopencl.enqueue_copy(
                self._pyopencl_command_queue,
                src=buffer.data, dest=data,
                host_origin=(data_origin,), buffer_origin=(buffer_origin,),
                region=shape,
                host_pitches=data_strides[1:], buffer_pitches=buffer_strides[1:])

    def enqueue_write_buffer_rect(
            self, buffer, data, buffer_origin, data_origin, shape,
            buffer_strides, data_strides, blocking=True):
        """Copy a region of host memory to a buffer. This is a low-level
        interface that ignores the shape, strides etc of the buffers, and
        treats them as byte arrays. It also only supports 3 or fewer dimensions. Use
        :py:meth:`~katsdpsigproc.accel.DeviceArray.set_region` for a high-level
        interface.

        Parameters
        ----------
        buffer : :class`pyopencl.array.Array`
            Target
        data : array-like
            Source
        buffer_origin, data_origin : int
            Offsets for the start of the copy, in bytes
        shape : sequence of int
            Shape of the region to copy (1-3 elements). The first dimension is
            a byte count.
        buffer_strides,data_strides : sequence of int
            Strides for the destination and source memory layout, with the same
            length as `shape`. The first element of each must be 1, and each
            element must be a factor of the next element.
        blocking : bool, optional
            If true, block until the transfer is complete.
        """
        assert buffer_strides[0] == 1
        assert data_strides[0] == 1
        if hasattr(data, 'enqueue_write_buffer_rect'):
            data.enqueue_write_buffer_rect(self, buffer, buffer_origin, data_origin, shape,
                                           buffer_strides, data_strides, blocking)
        else:
            pyopencl.enqueue_copy(
                self._pyopencl_command_queue,
                src=data, dest=buffer.data,
                host_origin=(data_origin,), buffer_origin=(buffer_origin,),
                region=shape,
                host_pitches=data_strides[1:], buffer_pitches=buffer_strides[1:])

    def enqueue_zero_buffer(self, buffer):
        """Fill a buffer with zero bytes.

        Parameters
        ----------
        buffer : :class:`pyopencl.array.Array`
        """
        pyopencl.enqueue_fill_buffer(self._pyopencl_command_queue, buffer.data,
                                     np.uint8(0), 0, buffer.data.get_info(pyopencl.mem_info.SIZE))

    @classmethod
    def _raw_arg(cls, arg):
        if isinstance(arg, pyopencl.array.Array):
            return arg.data
        else:
            return arg

    def _enqueue_kernel(self, kernel, args, global_size, local_size):
        """Enqueue a kernel to the command queue.

        This version returns the OpenCL event object. Refer to
        :meth:`enqueue_queue` for the public interface.
        """
        # PyOpenCL doesn't allow Array objects to be passed through
        args = [self._raw_arg(x) for x in args]
        # OpenCL allows local_size to be None, but this is not portable to
        # CUDA
        assert local_size is not None
        assert len(local_size) == len(global_size)
        return kernel._pyopencl_kernel(self._pyopencl_command_queue,
                                       global_size, local_size, *args)

    def enqueue_kernel(self, kernel, args, global_size, local_size):
        """Enqueue a kernel to the command queue.

        .. warning:: It is not thread-safe to call this function in two threads
            on the same kernel at the same time.

        Parameters
        ----------
        kernel : :class:`Kernel`
            Kernel to run
        args : sequence
            Arguments to pass to the kernel. Refer to the PyOpenCL
            documentation for details. Additionally, this function allows
            a :class:`pyopencl.array.Array` to be passed.
        global_size : tuple
            Number of work-items in each global dimension
        local_size : tuple
            Number of work-items in each local dimension. Must divide
            exactly into `global_size`.
        """
        self._enqueue_kernel(kernel, args, global_size, local_size)

    def enqueue_marker(self):
        """Create an event at this point in the command queue"""
        return Event(pyopencl.enqueue_marker(self._pyopencl_command_queue))

    def enqueue_wait_for_events(self, events):
        """Enqueue a barrier to wait for all events in `events`."""
        # OpenCL has some odd semantics for an empty wait list, hence the check
        if events:
            pyopencl.enqueue_barrier(self._pyopencl_command_queue,
                                     [x._pyopencl_event for x in events])

    def flush(self):
        """Start enqueued work running, but do not wait for it to complete"""
        self._pyopencl_command_queue.flush()

    def finish(self):
        """Block until all enqueued work has completed"""
        self._pyopencl_command_queue.finish()


class TuningCommandQueue(CommandQueue):
    """Command queue with extra facilities for autotuning. It keeps
    track of kernels that are enqueued since the last call to
    :meth:`start_tuning`, and reports the total time they consume
    when :meth:`stop_tuning` is called.
    """

    def __init__(self, *args, **kwargs):
        kwargs['profile'] = True
        super(TuningCommandQueue, self).__init__(*args, **kwargs)
        self.is_tuning = False
        self.events = []

    def start_tuning(self):
        self.is_tuning = True
        self.events = []

    def enqueue_kernel(self, kernel, args, global_size, local_size):
        if not self.is_tuning:
            super(TuningCommandQueue, self).enqueue_kernel(kernel, args, global_size, local_size)
        else:
            event = self._enqueue_kernel(kernel, args, global_size, local_size)
            self.events.append(event)

    def stop_tuning(self):
        self.finish()
        elapsed = 0.0
        if self.events:
            start = self.events[0].profile.start
            end = self.events[-1].profile.end
            elapsed = (end - start) * 1e-9
        self.is_tuning = False
        self.events = []
        return elapsed
