# Source code for arrayfire.interop

    #######################################################
    # Copyright (c) 2015, ArrayFire
    # All rights reserved.
    #
    # This file is distributed under 3-clause BSD license.
    # The complete license agreement can be obtained at:
    # http://arrayfire.com/licenses/BSD-3-Clause
    ########################################################

    """
    Interop with other python packages.

    This module provides interoperability with the following python packages.

         1. numpy
         2. pycuda
         3. pyopencl
    """

    from .array import *
    from .device import *

    try:
        import numpy as np
        from numpy import ndarray as NumpyArray
        from .data import reorder

        AF_NUMPY_FOUND=True

        def np_to_af_array(np_arr):
            """
            Convert numpy.ndarray to arrayfire.Array.

            Parameters
            ----------
            np_arr  : numpy.ndarray()

            Returns
            ---------
            af_arr  : arrayfire.Array()
            """

            in_shape = np_arr.shape
            in_ptr = np_arr.ctypes.data_as(ct.c_void_p)
            in_dtype = np_arr.dtype.char

            if (np_arr.flags['F_CONTIGUOUS']):
                return Array(in_ptr, in_shape, in_dtype)
            elif (np_arr.flags['C_CONTIGUOUS']):
                if np_arr.ndim == 1:
                    return Array(in_ptr, in_shape, in_dtype)
                elif np_arr.ndim == 2:
                    shape = (in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype)
                    return reorder(res, 1, 0)
                elif np_arr.ndim == 3:
                    shape = (in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype)
                    return reorder(res, 2, 1, 0)
                elif np_arr.ndim == 4:
                    shape = (in_shape[3], in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype)
                    return reorder(res, 3, 2, 1, 0)
                else:
                    raise RuntimeError("Unsupported ndim")
            else:
                return np_to_af_array(np.asfortranarray(np_arr))

        from_ndarray = np_to_af_array
    except:
        AF_NUMPY_FOUND=False

    try:
        import pycuda.gpuarray
        from pycuda.gpuarray import GPUArray as CudaArray
        AF_PYCUDA_FOUND=True

        def pycuda_to_af_array(pycu_arr):
            """
            Convert pycuda.gpuarray to arrayfire.Array

            Parameters
            -----------
            pycu_arr  : pycuda.GPUArray()

            Returns
            ----------
            af_arr    : arrayfire.Array()

            Note
            ----------
            The input array is copied to af.Array
            """

            in_ptr = pycu_arr.ptr
            in_shape = pycu_arr.shape
            in_dtype = pycu_arr.dtype.char

            if (pycu_arr.flags.f_contiguous):
                res = Array(in_ptr, in_shape, in_dtype, is_device=True)
                lock_array(res)
                res = res.copy()
                return res
            elif (pycu_arr.flags.c_contiguous):
                if pycu_arr.ndim == 1:
                    return Array(in_ptr, in_shape, in_dtype, is_device=True)
                elif pycu_arr.ndim == 2:
                    shape = (in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 1, 0)
                elif pycu_arr.ndim == 3:
                    shape = (in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 2, 1, 0)
                elif pycu_arr.ndim == 4:
                    shape = (in_shape[3], in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 3, 2, 1, 0)
                else:
                    raise RuntimeError("Unsupported ndim")
            else:
                return pycuda_to_af_array(pycu_arr.copy())
    except:
        AF_PYCUDA_FOUND=False

    try:
        from pyopencl.array import Array as OpenclArray
        from .opencl import add_device_context as _add_device_context
        from .opencl import set_device_context as _set_device_context
        from .opencl import get_device_id as _get_device_id
        from .opencl import get_context as _get_context
        AF_PYOPENCL_FOUND=True

        def pyopencl_to_af_array(pycl_arr):
            """
            Convert pyopencl.gpuarray to arrayfire.Array

            Parameters
            -----------
            pycl_arr  : pyopencl.Array()

            Returns
            ----------
            af_arr    : arrayfire.Array()

            Note
            ----------
            The input array is copied to af.Array
            """

            ctx = pycl_arr.context.int_ptr
            que = pycl_arr.queue.int_ptr
            dev = pycl_arr.queue.device.int_ptr

            dev_idx = None
            ctx_idx = None
            for n in range(get_device_count()):
                set_device(n)
                dev_idx = _get_device_id()
                ctx_idx = _get_context()
                if (dev_idx == dev and ctx_idx == ctx):
                    break

            if (dev_idx == None or ctx_idx == None or
                dev_idx != dev or ctx_idx != ctx):
                _add_device_context(dev, ctx, que)
                _set_device_context(dev, ctx)

            in_ptr = pycl_arr.base_data.int_ptr
            in_shape = pycl_arr.shape
            in_dtype = pycl_arr.dtype.char

            if (pycl_arr.flags.f_contiguous):
                res = Array(in_ptr, in_shape, in_dtype, is_device=True)
                lock_array(res)
                return res
            elif (pycl_arr.flags.c_contiguous):
                if pycl_arr.ndim == 1:
                    return Array(in_ptr, in_shape, in_dtype, is_device=True)
                elif pycl_arr.ndim == 2:
                    shape = (in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 1, 0)
                elif pycl_arr.ndim == 3:
                    shape = (in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 2, 1, 0)
                elif pycl_arr.ndim == 4:
                    shape = (in_shape[3], in_shape[2], in_shape[1], in_shape[0])
                    res = Array(in_ptr, shape, in_dtype, is_device=True)
                    lock_array(res)
                    return reorder(res, 3, 2, 1, 0)
                else:
                    raise RuntimeError("Unsupported ndim")
            else:
                return pyopencl_to_af_array(pycl_arr.copy())
    except:
        AF_PYOPENCL_FOUND=False

    def to_array(in_array):
        """
        Helper function to convert input from a different module to af.Array

        Parameters
        -------------

        in_array : array like object
                 Can be one of numpy.ndarray, pycuda.GPUArray, pyopencl.Array, array.array, list

        Returns
        --------------
        af.Array of same dimensions as input after copying the data from the input

        """
        if AF_NUMPY_FOUND and isinstance(in_array, NumpyArray):
            return np_to_af_array(in_array)
        if AF_PYCUDA_FOUND and isinstance(in_array, CudaArray):
            return pycuda_to_af_array(in_array)
        if AF_PYOPENCL_FOUND and isinstance(in_array, OpenclArray):
            return pyopencl_to_af_array(in_array)
        return Array(src=in_array)