from __future__ import annotations


__copyright__ = """
Copyright 2008-2021 Andreas Kloeckner
Copyright 2021 NVIDIA Corporation
"""

import os

import numpy as np


# {{{ add cuda lib dir to Python DLL path


def _search_on_path(filenames):
    """Find file on system path."""
    # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/52224

    from os import environ, pathsep
    from os.path import abspath, exists, join

    search_path = environ["PATH"]

    paths = search_path.split(pathsep)
    for path in paths:
        for filename in filenames:
            if exists(join(path, filename)):
                return abspath(join(path, filename))


def _add_cuda_libdir_to_dll_path():
    from os.path import dirname, exists, join

    cuda_path = os.environ.get("CUDA_PATH")

    if cuda_path is not None:
        # CUDA 13.0+ use bin/x64/ as dir
        lib_dir = join(cuda_path, "bin", "x64")
        if exists(lib_dir):
            os.add_dll_directory(lib_dir)
        else:
            os.add_dll_directory(join(cuda_path, "bin"))
        return

    nvcc_path = _search_on_path(["nvcc.exe"])
    if nvcc_path is not None:
        os.add_dll_directory(dirname(nvcc_path))
        return

    from warnings import warn

    warn(
        "Unable to discover CUDA installation directory "
        "while attempting to add it to Python's DLL path. "
        "Either set the 'CUDA_PATH' environment variable "
        "or ensure that 'nvcc.exe' is on the path.", stacklevel=2
    )


try:
    os.add_dll_directory  # noqa: B018
except AttributeError:
    # likely not on Py3.8 and Windows
    # https://github.com/inducer/pycuda/issues/213
    pass
else:
    _add_cuda_libdir_to_dll_path()

# }}}


try:
    from pycuda._driver import *  # noqa
except ImportError as e:
    if "_v2" in str(e):
        from warnings import warn

        warn(
            "Failed to import the CUDA driver interface, with an error "
            "message indicating that the version of your CUDA header "
            "does not match the version of your CUDA driver.", stacklevel=2
        )
    raise


_memoryview = memoryview
_my_bytes = bytes


try:
    ManagedAllocationOrStub = ManagedAllocation
except NameError:
    # Provide ManagedAllocationOrStub if not on CUDA 6.
    # This avoids having to do a version check in a high-traffic code path below.

    class ManagedAllocationOrStub:
        pass


CUDA_DEBUGGING = False


def set_debugging(flag=True):
    global CUDA_DEBUGGING
    CUDA_DEBUGGING = flag


class CompileError(Error):
    def __init__(self, msg, command_line, stdout=None, stderr=None):
        self.msg = msg
        self.command_line = command_line
        self.stdout = stdout
        self.stderr = stderr

    def __str__(self):
        result = self.msg
        if self.command_line:
            try:
                result += "\n[command: %s]" % (" ".join(self.command_line))
            except Exception as e:
                print(e)
        if self.stdout:
            result += "\n[stdout:\n%s]" % self.stdout
        if self.stderr:
            result += "\n[stderr:\n%s]" % self.stderr

        return result


class ArgumentHandler:
    def __init__(self, ary):
        self.array = ary
        self.dev_alloc = None

    def get_device_alloc(self):
        if self.dev_alloc is None:
            try:
                self.dev_alloc = mem_alloc_like(self.array)
            except AttributeError:
                raise TypeError(
                    "could not determine array length of '%s': "
                    "unsupported array type or not an array"
                    % type(self.array)
                )
        return self.dev_alloc

    def pre_call(self, stream):
        pass


class In(ArgumentHandler):
    def pre_call(self, stream):
        if stream is not None:
            memcpy_htod(self.get_device_alloc(), self.array)
        else:
            memcpy_htod(self.get_device_alloc(), self.array)


class Out(ArgumentHandler):
    def post_call(self, stream):
        if stream is not None:
            memcpy_dtoh(self.array, self.get_device_alloc())
        else:
            memcpy_dtoh(self.array, self.get_device_alloc())


class InOut(In, Out):
    pass


def _add_functionality():
    def device_get_attributes(dev):
        result = {}

        for att_name in dir(device_attribute):
            if not att_name[0].isupper():
                continue

            att_id = getattr(device_attribute, att_name)

            try:
                att_value = dev.get_attribute(att_id)
            except LogicError as e:
                from warnings import warn

                warn(
                    "CUDA driver raised '%s' when querying '%s' on '%s'"
                    % (e, att_name, dev), stacklevel=2
                )
            else:
                result[att_id] = att_value

        return result

    def device___getattr__(dev, name):
        return dev.get_attribute(getattr(device_attribute, name.upper()))

    def _build_arg_buf(args):
        handlers = []

        arg_data = []
        format = ""
        for i, arg in enumerate(args):
            if isinstance(arg, np.number):
                arg_data.append(arg)
                format += arg.dtype.char
            elif isinstance(arg, (DeviceAllocation, PooledDeviceAllocation)):
                arg_data.append(int(arg))
                format += "P"
            elif isinstance(arg, ArgumentHandler):
                handlers.append(arg)
                arg_data.append(int(arg.get_device_alloc()))
                format += "P"
            elif isinstance(arg, np.ndarray):
                if isinstance(arg.base, ManagedAllocationOrStub):
                    arg_data.append(int(arg.base))
                    format += "P"
                else:
                    arg_data.append(arg)
                    format += "%ds" % arg.nbytes
            elif isinstance(arg, np.void):
                arg_data.append(_my_bytes(_memoryview(arg)))
                format += "%ds" % arg.itemsize
            else:
                cai = getattr(arg, "__cuda_array_interface__", None)
                if cai:
                    arg_data.append(cai["data"][0])
                    format += "P"
                    continue

                try:
                    gpudata = np.uintp(arg.gpudata)
                except AttributeError:
                    raise TypeError("invalid type on parameter #%d (0-based)" % i)
                else:
                    # for gpuarrays
                    arg_data.append(int(gpudata))
                    format += "P"

        from pycuda._pvt_struct import pack

        return handlers, pack(format, *arg_data)

    # {{{ pre-CUDA 4 call interface (stateful)

    def function_param_set_pre_v4(func, *args):
        handlers = []

        handlers, buf = _build_arg_buf(args)

        func._param_setv(0, buf)
        func._param_set_size(len(buf))

        return handlers

    def function_call_pre_v4(func, *args, **kwargs):
        grid = kwargs.pop("grid", (1, 1))
        stream = kwargs.pop("stream", None)
        block = kwargs.pop("block", None)
        shared = kwargs.pop("shared", None)
        texrefs = kwargs.pop("texrefs", [])
        time_kernel = kwargs.pop("time_kernel", False)

        if kwargs:
            raise ValueError(
                "extra keyword arguments: %s" % (",".join(kwargs.keys()))
            )

        if block is None:
            raise ValueError("must specify block size")

        func._set_block_shape(*block)
        handlers = func._param_set(*args)
        if shared is not None:
            func._set_shared_size(shared)

        for handler in handlers:
            handler.pre_call(stream)

        for texref in texrefs:
            func.param_set_texref(texref)

        post_handlers = [
            handler for handler in handlers if hasattr(handler, "post_call")
        ]

        if stream is None:
            if time_kernel:
                Context.synchronize()

                from time import time

                start_time = time()
            func._launch_grid(*grid)
            if post_handlers or time_kernel:
                Context.synchronize()

                if time_kernel:
                    run_time = time() - start_time

                for handler in post_handlers:
                    handler.post_call(stream)

                if time_kernel:
                    return run_time
        else:
            assert (
                not time_kernel
            ), "Can't time the kernel on an asynchronous invocation"
            func._launch_grid_async(grid[0], grid[1], stream)

            if post_handlers:
                for handler in post_handlers:
                    handler.post_call(stream)

    def function_prepare_pre_v4(func, arg_types, block=None, shared=None, texrefs=None):
        from warnings import warn

        if texrefs is None:
            texrefs = []
        if block is not None:
            warn(
                "setting the block size in Function.prepare is deprecated",
                DeprecationWarning,
                stacklevel=2,
            )
            func._set_block_shape(*block)

        if shared is not None:
            warn(
                "setting the shared memory size in Function.prepare is deprecated",
                DeprecationWarning,
                stacklevel=2,
            )
            func._set_shared_size(shared)

        func.texrefs = texrefs

        func.arg_format = ""

        for _i, arg_type in enumerate(arg_types):
            if (
                isinstance(arg_type, type)
                and np is not None
                and np.number in arg_type.__mro__
            ):
                func.arg_format += np.dtype(arg_type).char
            elif isinstance(arg_type, str):
                func.arg_format += arg_type
            else:
                func.arg_format += np.dtype(np.uintp).char

        from pycuda._pvt_struct import calcsize

        func._param_set_size(calcsize(func.arg_format))

        return func

    def function_prepared_call_pre_v4(func, grid, block, *args, **kwargs):
        if isinstance(block, tuple):
            func._set_block_shape(*block)
        else:
            from warnings import warn

            warn(
                "Not passing the block size to prepared_call is deprecated as of "
                "version 2011.1.",
                DeprecationWarning,
                stacklevel=2,
            )
            args = (block, *args)

        shared_size = kwargs.pop("shared_size", None)
        if shared_size is not None:
            func._set_shared_size(shared_size)

        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        func._param_setv(0, pack(func.arg_format, *args))

        for texref in func.texrefs:
            func.param_set_texref(texref)

        func._launch_grid(*grid)

    def function_prepared_timed_call_pre_v4(func, grid, block, *args, **kwargs):
        if isinstance(block, tuple):
            func._set_block_shape(*block)
        else:
            from warnings import warn

            warn(
                "Not passing the block size to prepared_timed_call is "
                "deprecated as of version 2011.1.",
                DeprecationWarning,
                stacklevel=2,
            )
            args = (block, *args)

        shared_size = kwargs.pop("shared_size", None)
        if shared_size is not None:
            func._set_shared_size(shared_size)

        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        func._param_setv(0, pack(func.arg_format, *args))

        for texref in func.texrefs:
            func.param_set_texref(texref)

        start = Event()
        end = Event()

        start.record()
        func._launch_grid(*grid)
        end.record()

        def get_call_time():
            end.synchronize()
            return end.time_since(start) * 1e-3

        return get_call_time

    def function_prepared_async_call_pre_v4(func, grid, block, stream, *args, **kwargs):
        if isinstance(block, tuple):
            func._set_block_shape(*block)
        else:
            from warnings import warn

            warn(
                "Not passing the block size to prepared_async_call is "
                "deprecated as of version 2011.1.",
                DeprecationWarning,
                stacklevel=2,
            )
            args = (stream, *args)
            stream = block

        shared_size = kwargs.pop("shared_size", None)
        if shared_size is not None:
            func._set_shared_size(shared_size)

        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        func._param_setv(0, pack(func.arg_format, *args))

        for texref in func.texrefs:
            func.param_set_texref(texref)

        if stream is None:
            func._launch_grid(*grid)
        else:
            grid_x, grid_y = grid
            func._launch_grid_async(grid_x, grid_y, stream)

    # }}}

    # {{{ CUDA 4+ call interface (stateless)

    def function_call(func, *args, **kwargs):
        grid = kwargs.pop("grid", (1, 1))
        stream = kwargs.pop("stream", None)
        block = kwargs.pop("block", None)
        shared = kwargs.pop("shared", 0)
        texrefs = kwargs.pop("texrefs", [])
        time_kernel = kwargs.pop("time_kernel", False)

        if kwargs:
            raise ValueError(
                "extra keyword arguments: %s" % (",".join(kwargs.keys()))
            )

        if block is None:
            raise ValueError("must specify block size")

        func._set_block_shape(*block)
        handlers, arg_buf = _build_arg_buf(args)

        for handler in handlers:
            handler.pre_call(stream)

        for texref in texrefs:
            func.param_set_texref(texref)

        post_handlers = [
            handler for handler in handlers if hasattr(handler, "post_call")
        ]

        if stream is None:
            if time_kernel:
                Context.synchronize()

                from time import time

                start_time = time()

            func._launch_kernel(grid, block, arg_buf, shared, None)

            if post_handlers or time_kernel:
                Context.synchronize()

                if time_kernel:
                    run_time = time() - start_time

                for handler in post_handlers:
                    handler.post_call(stream)

                if time_kernel:
                    return run_time
        else:
            assert (
                not time_kernel
            ), "Can't time the kernel on an asynchronous invocation"
            func._launch_kernel(grid, block, arg_buf, shared, stream)

            if post_handlers:
                for handler in post_handlers:
                    handler.post_call(stream)

    def function_prepare(func, arg_types, texrefs=None):
        if texrefs is None:
            texrefs = []
        func.texrefs = texrefs

        func.arg_format = ""

        for _i, arg_type in enumerate(arg_types):
            if isinstance(arg_type, type) and np.number in arg_type.__mro__:
                func.arg_format += np.dtype(arg_type).char
            elif isinstance(arg_type, np.dtype):
                if arg_type.char == "V":
                    func.arg_format += "%ds" % arg_type.itemsize
                else:
                    func.arg_format += arg_type.char
            elif isinstance(arg_type, str):
                func.arg_format += arg_type
            else:
                func.arg_format += np.dtype(np.uintp).char

        return func

    def function_prepared_call(func, grid, block, *args, **kwargs):
        if isinstance(block, tuple):
            func._set_block_shape(*block)
        else:
            from warnings import warn

            warn(
                "Not passing the block size to prepared_call is deprecated as of "
                "version 2011.1.",
                DeprecationWarning,
                stacklevel=2,
            )
            args = (block, *args)

        shared_size = kwargs.pop("shared_size", 0)

        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        arg_buf = pack(func.arg_format, *args)

        for texref in func.texrefs:
            func.param_set_texref(texref)

        func._launch_kernel(grid, block, arg_buf, shared_size, None)

    def function_prepared_timed_call(func, grid, block, *args, **kwargs):
        shared_size = kwargs.pop("shared_size", 0)
        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        arg_buf = pack(func.arg_format, *args)

        for texref in func.texrefs:
            func.param_set_texref(texref)

        start = Event()
        end = Event()

        start.record()
        func._launch_kernel(grid, block, arg_buf, shared_size, None)
        end.record()

        def get_call_time():
            end.synchronize()
            return end.time_since(start) * 1e-3

        return get_call_time

    def function_prepared_async_call(func, grid, block, stream, *args, **kwargs):
        if isinstance(block, tuple):
            func._set_block_shape(*block)
        else:
            from warnings import warn

            warn(
                "Not passing the block size to prepared_async_call is "
                "deprecated as of version 2011.1.",
                DeprecationWarning,
                stacklevel=2,
            )
            args = (stream, *args)
            stream = block

        shared_size = kwargs.pop("shared_size", 0)

        if kwargs:
            raise TypeError(
                "unknown keyword arguments: " + ", ".join(kwargs.keys())
            )

        from pycuda._pvt_struct import pack

        arg_buf = pack(func.arg_format, *args)

        for texref in func.texrefs:
            func.param_set_texref(texref)

        func._launch_kernel(grid, block, arg_buf, shared_size, stream)

    # }}}

    def function___getattr__(self, name):
        if get_version() >= (2, 2):
            return self.get_attribute(getattr(function_attribute, name.upper()))
        else:
            if name == "num_regs":
                return self._hacky_registers
            elif name == "shared_size_bytes":
                return self._hacky_smem
            elif name == "local_size_bytes":
                return self._hacky_lmem
            else:
                raise AttributeError("no attribute '%s' in Function" % name)

    def mark_func_method_deprecated(func):
        def new_func(*args, **kwargs):
            from warnings import warn

            warn(
                "'%s' has been deprecated in version 2011.1. Please use "
                "the stateless launch interface instead." % func.__name__[1:],
                DeprecationWarning,
                stacklevel=2,
            )
            return func(*args, **kwargs)

        from functools import update_wrapper
        update_wrapper(new_func, func)

        return new_func

    Device.get_attributes = device_get_attributes
    Device.__getattr__ = device___getattr__

    if get_version() >= (4,):
        Function.__call__ = function_call
        Function.prepare = function_prepare
        Function.prepared_call = function_prepared_call
        Function.prepared_timed_call = function_prepared_timed_call
        Function.prepared_async_call = function_prepared_async_call
    else:
        Function._param_set = function_param_set_pre_v4
        Function.__call__ = function_call_pre_v4
        Function.prepare = function_prepare_pre_v4
        Function.prepared_call = function_prepared_call_pre_v4
        Function.prepared_timed_call = function_prepared_timed_call_pre_v4
        Function.prepared_async_call = function_prepared_async_call_pre_v4

        for meth_name in [
            "set_block_shape",
            "set_shared_size",
            "param_set_size",
            "param_set",
            "param_seti",
            "param_setf",
            "param_setv",
            "launch",
            "launch_grid",
            "launch_grid_async",
        ]:
            setattr(
                Function,
                meth_name,
                mark_func_method_deprecated(getattr(Function, "_" + meth_name)),
            )

    Function.__getattr__ = function___getattr__


_add_functionality()


# {{{ pagelocked numpy arrays


def pagelocked_zeros(shape, dtype, order="C", mem_flags=0):
    result = pagelocked_empty(shape, dtype, order, mem_flags)
    result.fill(0)
    return result


def pagelocked_empty_like(array, mem_flags=0):
    if array.flags.c_contiguous:
        order = "C"
    elif array.flags.f_contiguous:
        order = "F"
    else:
        raise ValueError("could not detect array order")

    return pagelocked_empty(array.shape, array.dtype, order, mem_flags)


def pagelocked_zeros_like(array, mem_flags=0):
    result = pagelocked_empty_like(array, mem_flags)
    result.fill(0)
    return result


# }}}


# {{{ aligned numpy arrays


def aligned_zeros(shape, dtype, order="C", alignment=4096):
    result = aligned_empty(shape, dtype, order, alignment)
    result.fill(0)
    return result


def aligned_empty_like(array, alignment=4096):
    if array.flags.c_contiguous:
        order = "C"
    elif array.flags.f_contiguous:
        order = "F"
    else:
        raise ValueError("could not detect array order")

    return aligned_empty(array.shape, array.dtype, order, alignment)


def aligned_zeros_like(array, alignment=4096):
    result = aligned_empty_like(array, alignment)
    result.fill(0)
    return result


# }}}


# {{{ managed numpy arrays (CUDA Unified Memory)


def managed_zeros(shape, dtype, order="C", mem_flags=0):
    result = managed_empty(shape, dtype, order, mem_flags)
    result.fill(0)
    return result


def managed_empty_like(array, mem_flags=0):
    if array.flags.c_contiguous:
        order = "C"
    elif array.flags.f_contiguous:
        order = "F"
    else:
        raise ValueError("could not detect array order")

    return managed_empty(array.shape, array.dtype, order, mem_flags)


def managed_zeros_like(array, mem_flags=0):
    result = managed_empty_like(array, mem_flags)
    result.fill(0)
    return result


# }}}


def mem_alloc_like(ary):
    return mem_alloc(ary.nbytes)


# {{{ array handling


def dtype_to_array_format(dtype):
    if dtype == np.uint8:
        return array_format.UNSIGNED_INT8
    elif dtype == np.uint16:
        return array_format.UNSIGNED_INT16
    elif dtype == np.uint32:
        return array_format.UNSIGNED_INT32
    elif dtype == np.int8:
        return array_format.SIGNED_INT8
    elif dtype == np.int16:
        return array_format.SIGNED_INT16
    elif dtype == np.int32:
        return array_format.SIGNED_INT32
    elif dtype == np.float32:
        return array_format.FLOAT
    else:
        raise TypeError("cannot convert dtype '%s' to array format" % dtype)


def matrix_to_array(matrix, order, allow_double_hack=False):
    if order.upper() == "C":
        h, w = matrix.shape
        stride = 0
    elif order.upper() == "F":
        w, h = matrix.shape
        stride = -1
    else:
        raise LogicError("order must be either F or C")

    matrix = np.asarray(matrix, order=order)
    descr = ArrayDescriptor()

    descr.width = w
    descr.height = h

    if matrix.dtype == np.float64 and allow_double_hack:
        descr.format = array_format.SIGNED_INT32
        descr.num_channels = 2
    else:
        descr.format = dtype_to_array_format(matrix.dtype)
        descr.num_channels = 1

    ary = Array(descr)

    copy = Memcpy2D()
    copy.set_src_host(matrix)
    copy.set_dst_array(ary)
    copy.width_in_bytes = copy.src_pitch = copy.dst_pitch = matrix.strides[stride]
    copy.height = h
    copy(aligned=True)

    return ary


def np_to_array(nparray, order, allowSurfaceBind=False):  # noqa: N803
    case = order in ["C", "F"]
    if not case:
        raise LogicError("order must be either F or C")

    dimension = len(nparray.shape)
    if dimension == 2:
        if order == "C":
            stride = 0
        if order == "F":
            stride = -1
        h, w = nparray.shape
        d = 1
        if allowSurfaceBind:
            descrArr = ArrayDescriptor3D()
            descrArr.width = w
            descrArr.height = h
            descrArr.depth = d
        else:
            descrArr = ArrayDescriptor()
            descrArr.width = w
            descrArr.height = h
    elif dimension == 3:
        if order == "C":
            stride = 1
        if order == "F":
            stride = 1
        d, h, w = nparray.shape
        descrArr = ArrayDescriptor3D()
        descrArr.width = w
        descrArr.height = h
        descrArr.depth = d
    else:
        raise LogicError(
            "CUDArrays dimensions 2 or 3 supported in CUDA at the moment ... "
        )

    if nparray.dtype == np.complex64:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int2 (hi=re,lo=im) structure
        descrArr.num_channels = 2
    elif nparray.dtype == np.float64:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int2 (hi,lo) structure
        descrArr.num_channels = 2
    elif nparray.dtype == np.complex128:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int4 (re=(hi,lo),im=(hi,lo)) structure
        descrArr.num_channels = 4
    else:
        descrArr.format = dtype_to_array_format(nparray.dtype)
        descrArr.num_channels = 1

    if allowSurfaceBind:
        if dimension == 2:
            descrArr.flags |= array3d_flags.ARRAY3D_LAYERED
        descrArr.flags |= array3d_flags.SURFACE_LDST

    cudaArray = Array(descrArr)
    if allowSurfaceBind or dimension == 3:
        copy3D = Memcpy3D()
        copy3D.set_src_host(nparray)
        copy3D.set_dst_array(cudaArray)
        copy3D.width_in_bytes = copy3D.src_pitch = nparray.strides[stride]
        copy3D.src_height = copy3D.height = h
        copy3D.depth = d
        copy3D()
        return cudaArray
    else:
        copy2D = Memcpy2D()
        copy2D.set_src_host(nparray)
        copy2D.set_dst_array(cudaArray)
        copy2D.width_in_bytes = copy2D.src_pitch = nparray.strides[stride]
        copy2D.src_height = copy2D.height = h
        copy2D(aligned=True)
        return cudaArray


def gpuarray_to_array(gpuarray, order, allowSurfaceBind=False):  # noqa: N803
    case = order in ["C", "F"]
    if not case:
        raise LogicError("order must be either F or C")

    dimension = len(gpuarray.shape)
    if dimension == 2:
        if order == "C":
            stride = 0
        if order == "F":
            stride = -1
        h, w = gpuarray.shape
        d = 1
        if allowSurfaceBind:
            descrArr = ArrayDescriptor3D()
            descrArr.width = int(w)
            descrArr.height = int(h)
            descrArr.depth = int(d)
        else:
            descrArr = ArrayDescriptor()
            descrArr.width = int(w)
            descrArr.height = int(h)
    elif dimension == 3:
        if order == "C":
            stride = 1
        if order == "F":
            stride = 1
        d, h, w = gpuarray.shape
        descrArr = ArrayDescriptor3D()
        descrArr.width = int(w)
        descrArr.height = int(h)
        descrArr.depth = int(d)
    else:
        raise LogicError(
            "CUDArray dimensions 2 and 3 supported in CUDA at the moment ... "
        )

    if gpuarray.dtype == np.complex64:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int2 (hi=re,lo=im) structure
        descrArr.num_channels = 2
    elif gpuarray.dtype == np.float64:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int2 (hi,lo) structure
        descrArr.num_channels = 2
    elif gpuarray.dtype == np.complex128:
        descrArr.format = (
            array_format.SIGNED_INT32
        )  # Reading data as int4 (re=(hi,lo),im=(hi,lo)) structure
        descrArr.num_channels = 4
    else:
        descrArr.format = dtype_to_array_format(gpuarray.dtype)
        descrArr.num_channels = 1

    if allowSurfaceBind:
        if dimension == 2:
            descrArr.flags |= array3d_flags.ARRAY3D_LAYERED
        descrArr.flags |= array3d_flags.SURFACE_LDST

    cudaArray = Array(descrArr)
    if allowSurfaceBind or dimension == 3:
        copy3D = Memcpy3D()
        copy3D.set_src_device(gpuarray.ptr)
        copy3D.set_dst_array(cudaArray)
        copy3D.width_in_bytes = copy3D.src_pitch = gpuarray.strides[stride]
        copy3D.src_height = copy3D.height = int(h)
        copy3D.depth = int(d)
        copy3D()
        return cudaArray
    else:
        copy2D = Memcpy2D()
        copy2D.set_src_device(gpuarray.ptr)
        copy2D.set_dst_array(cudaArray)
        copy2D.width_in_bytes = copy2D.src_pitch = gpuarray.strides[stride]
        copy2D.src_height = copy2D.height = int(h)
        copy2D(aligned=True)
        return cudaArray


def make_multichannel_2d_array(ndarray, order):
    """Channel count has to be the first dimension of the C{ndarray}."""

    descr = ArrayDescriptor()

    if order.upper() == "C":
        h, w, num_channels = ndarray.shape
        stride = 0
    elif order.upper() == "F":
        num_channels, w, h = ndarray.shape
        stride = 2
    else:
        raise LogicError("order must be either F or C")

    descr.width = w
    descr.height = h
    descr.format = dtype_to_array_format(ndarray.dtype)
    descr.num_channels = num_channels

    ary = Array(descr)

    copy = Memcpy2D()
    copy.set_src_host(ndarray)
    copy.set_dst_array(ary)
    copy.width_in_bytes = copy.src_pitch = copy.dst_pitch = ndarray.strides[stride]
    copy.height = h
    copy(aligned=True)

    return ary


def bind_array_to_texref(ary, texref):
    texref.set_array(ary)
    texref.set_address_mode(0, address_mode.CLAMP)
    texref.set_address_mode(1, address_mode.CLAMP)
    texref.set_filter_mode(filter_mode.POINT)


# }}}


def matrix_to_texref(matrix, texref, order):
    bind_array_to_texref(matrix_to_array(matrix, order), texref)


# {{{ device copies


def to_device(bf_obj):

    bf = memoryview(bf_obj).tobytes()
    result = mem_alloc(len(bf))
    memcpy_htod(result, bf)
    return result


def from_device(devptr, shape, dtype, order="C"):
    result = np.empty(shape, dtype, order)
    memcpy_dtoh(result, devptr)
    return result


def from_device_like(devptr, other_ary):
    result = np.empty_like(other_ary)
    memcpy_dtoh(result, devptr)
    return result


# }}}

# vim: fdm=marker
