# -*- coding: utf-8 -*-

# PyVkFFT
#   (c) 2021- : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

__all__ = ['fftn', 'ifftn', 'rfftn', 'irfftn', 'dctn', 'idctn', 'dstn', 'idstn',
           'vkfft_version', 'clear_vkfftapp_cache',
           'has_pycuda', 'has_opencl', 'has_cupy']

from enum import Enum
from functools import lru_cache
import numpy as np
from .base import complex32, strides_nonzero
from . import config

try:
    from .cuda import VkFFTApp as VkFFTApp_cuda, has_pycuda, has_cupy, vkfft_version

    if has_pycuda:
        import pycuda.gpuarray as cua
        import pycuda.driver as cu_drv
    if has_cupy:
        import cupy as cp
except (ImportError, OSError):
    has_cupy, has_pycuda = False, False

try:
    from .opencl import VkFFTApp as VkFFTApp_cl, cla, vkfft_version

    has_opencl = True
except (ImportError, OSError):
    has_opencl = False


class Backend(Enum):
    """ Backend language & library"""
    UNKNOWN = 0
    PYCUDA = 1
    PYOPENCL = 2
    CUPY = 3


def _prepare_transform(src, dest, cl_queue, cuda_stream, r2c=False, r2c_odd=False):
    """
    Determine the backend from the input data.
    Create the destination array if necessary.

    :param src: the source GPU array
    :param dest: the destination array. If None, a new GPU array is created.
    :param cl_queue: the opencl queue to use, or None
    :param cuda_stream: the cuda stream to use (from cupy or pycuda), or None
    :param r2c: if True, this is for an R2C transform, so adapt the destination
        array accordingly.
    :param r2c_odd: True if the r2c transform has an odd size (for the real
        array) along the x-axis. This is only needed for a c2r out-of-place
        transform to determine the destination x-axis size.
    :return: a tuple (backend, inplace, dest, cl_queue, cuda_stream, devctx),
        also appending the destination dtype for a r2c transform.
        devctx is either the device or context unique ptr or id - this
        is only used to cache the VkFFTapp (e.g. making sure the app is
        re-instantiated if the cuda device changes). The cuda_stream
        output will be the int ptr.
    """
    backend = Backend.UNKNOWN
    # fast axis is the last only for C-ordered arrays
    # We don't use np.argmin(src.strides) because of special cases
    # with axis size=1
    fast_axis = src.ndim - 1 - np.argmin(strides_nonzero(src.strides)[::-1])
    if fast_axis == src.ndim - 1:
        order = 'C'
    else:
        order = 'F'
    if r2c:
        if src.dtype in [np.float16, np.float32, np.float64]:
            sh = list(src.shape)
            sh[fast_axis] = sh[fast_axis] // 2 + 1
            dtype = np.complex64
            if src.dtype == np.float16:
                dtype = complex32
            elif src.dtype == np.float64:
                dtype = np.complex128
        else:
            sh = list(src.shape)
            sh[fast_axis] = (sh[fast_axis] - 1) * 2
            if r2c_odd:
                # This is only true for an out-of-place transform,
                # but sh is only used if dest is None
                sh[fast_axis] += 1
            dtype = np.float32
            if src.dtype == complex32:
                dtype = np.float16
            elif src.dtype == np.complex128:
                dtype = np.float64
    else:
        sh, dtype = None, None
    devctx = None
    if has_pycuda:
        if isinstance(src, cua.GPUArray):
            backend = Backend.PYCUDA
            # Must cast the gpudata to int as it can either be a DeviceAllocation object
            # or an int (e.g. when using a view of another array)
            src_ptr = int(src.gpudata)
            if dest is None:
                if r2c:
                    dest = cua.empty(tuple(sh), dtype=dtype, allocator=src.allocator, order=order)
                else:
                    dest = cua.empty_like(src)
            dest_ptr = int(dest.gpudata)
            if cuda_stream is None:
                devctx = cu_drv.Context.get_current()
            elif isinstance(cuda_stream, cu_drv.Stream):
                # Pass an int to make sure it is hashable
                cuda_stream = cuda_stream.handle

    if backend == Backend.UNKNOWN and has_opencl:
        if isinstance(src, cla.Array):
            backend = Backend.PYOPENCL
            src_ptr = src.data.int_ptr
            if dest is None:
                if r2c:
                    dest = cla.empty(src.queue, tuple(sh), dtype=dtype, allocator=src.allocator, order=order)
                else:
                    dest = cla.empty_like(src)
            dest_ptr = dest.data.int_ptr
            if cl_queue is None:
                cl_queue = src.queue
            devctx = cl_queue.context

    if backend == Backend.UNKNOWN and has_cupy:
        if isinstance(src, cp.ndarray):
            backend = Backend.CUPY
            src_ptr = src.__cuda_array_interface__['data'][0]
            if dest is None:
                if r2c:
                    dest = cp.empty(tuple(sh), dtype=dtype, order=order)
                else:
                    dest = cp.empty_like(src)
            dest_ptr = dest.__cuda_array_interface__['data'][0]
            if cuda_stream is None:
                cuda_stream = cp.cuda.get_current_stream()
            if isinstance(cuda_stream, cp.cuda.Stream):
                cuda_stream = cuda_stream.ptr
            devctx = cp.cuda.Device().id

    if backend == Backend.UNKNOWN:
        raise RuntimeError("Could note determine the type of GPU array supplied, or the "
                           "corresponding backend is not installed "
                           "(has_pycuda=%d, has_pyopencl=%d, has_cupy=%d)" %
                           (has_pycuda, has_opencl, has_cupy))

    inplace = dest_ptr == src_ptr
    if r2c:
        if inplace:
            dest = src.view(dtype=dtype)
        return backend, inplace, dest, cl_queue, cuda_stream, devctx, dtype
    else:
        return backend, inplace, dest, cl_queue, cuda_stream, devctx


@lru_cache(maxsize=config.FFT_CACHE_NB)
def _get_fft_app(backend, shape, dtype, inplace, ndim, axes, norm, cuda_stream, cl_queue,
                 devctx, strides=None, tune=False):
    del devctx  # Variable is just used for proper lru_cache
    sback = {Backend.PYCUDA: 'pycuda', Backend.CUPY: 'cupy', Backend.PYOPENCL: 'pyopencl'}[backend]
    tune_config = {'backend': sback} if tune else None
    if backend in [Backend.PYCUDA, Backend.CUPY]:
        return VkFFTApp_cuda(shape, dtype, ndim=ndim, inplace=inplace,
                             stream=cuda_stream, norm=norm, axes=axes, strides=strides,
                             tune_config=tune_config)
    elif backend == Backend.PYOPENCL:
        return VkFFTApp_cl(shape, dtype, cl_queue, ndim=ndim, inplace=inplace,
                           norm=norm, axes=axes, strides=strides,
                           tune_config=tune_config)


@lru_cache(maxsize=config.FFT_CACHE_NB)
def _get_rfft_app(backend, shape, dtype, inplace, ndim, norm, cuda_stream, cl_queue,
                  devctx, strides=None, tune=False, r2c_odd=False):
    del devctx  # Variable is just used for proper lru_cache
    sback = {Backend.PYCUDA: 'pycuda', Backend.CUPY: 'cupy', Backend.PYOPENCL: 'pyopencl'}[backend]
    tune_config = {'backend': sback} if tune else None
    if backend in [Backend.PYCUDA, Backend.CUPY]:
        return VkFFTApp_cuda(shape, dtype, ndim=ndim, inplace=inplace,
                             stream=cuda_stream, norm=norm, r2c=True, strides=strides,
                             tune_config=tune_config, r2c_odd=r2c_odd)
    elif backend == Backend.PYOPENCL:
        return VkFFTApp_cl(shape, dtype, cl_queue, ndim=ndim, inplace=inplace,
                           norm=norm, r2c=True, strides=strides,
                           tune_config=tune_config, r2c_odd=r2c_odd)


@lru_cache(maxsize=config.FFT_CACHE_NB)
def _get_dct_app(backend, shape, dtype, inplace, ndim, norm, dct_type,
                 cuda_stream, cl_queue, devctx, tune=False):
    del devctx  # Variable is just used for proper lru_cache
    sback = {Backend.PYCUDA: 'pycuda', Backend.CUPY: 'cupy', Backend.PYOPENCL: 'pyopencl'}[backend]
    tune_config = {'backend': sback} if tune else None
    if backend in [Backend.PYCUDA, Backend.CUPY]:
        return VkFFTApp_cuda(shape, dtype, ndim=ndim, inplace=inplace,
                             stream=cuda_stream, norm=norm, dct=dct_type,
                             tune_config=tune_config)
    elif backend == Backend.PYOPENCL:
        return VkFFTApp_cl(shape, dtype, cl_queue, ndim=ndim, inplace=inplace,
                           norm=norm, dct=dct_type,
                           tune_config=tune_config)


@lru_cache(maxsize=config.FFT_CACHE_NB)
def _get_dst_app(backend, shape, dtype, inplace, ndim, norm, dst_type,
                 cuda_stream, cl_queue, devctx, tune=False):
    del devctx  # Variable is just used for proper lru_cache
    sback = {Backend.PYCUDA: 'pycuda', Backend.CUPY: 'cupy', Backend.PYOPENCL: 'pyopencl'}[backend]
    tune_config = {'backend': sback} if tune else None
    if backend in [Backend.PYCUDA, Backend.CUPY]:
        return VkFFTApp_cuda(shape, dtype, ndim=ndim, inplace=inplace,
                             stream=cuda_stream, norm=norm, dst=dst_type,
                             tune_config=tune_config)
    elif backend == Backend.PYOPENCL:
        return VkFFTApp_cl(shape, dtype, cl_queue, ndim=ndim, inplace=inplace,
                           norm=norm, dst=dst_type,
                           tune_config=tune_config)


def fftn(src, dest=None, ndim=None, norm=1, axes=None, cuda_stream=None, cl_queue=None,
         return_scale=False, tune=False):
    """
    Perform a FFT on a GPU array, automatically creating the VkFFTApp
    and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: if 0 (un-normalised), every transform multiplies the L2 norm
        of the array by the transform size.
        if 1 (the default) or "backward", the inverse transform divides the
        L2 norm by the array size, so FFT+iFFT will keep the array norm.
        if "ortho", each transform will keep the L2 norm, but that will
        involve an extra read & write operation.
    :param axes: a list or tuple of axes along which the transform is made.
        if None, the transform is done along the ndim fastest axes, or all
        axes if ndim is None. Not allowed for R2C transforms
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param return_scale: if True, return the scale factor by which the result
        must be multiplied to keep its L2 norm after the transform
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array if return_scale is False, or (dest, scale)
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_fft_app(backend, src.shape, src.dtype, inplace, ndim, axes, norm, cuda_stream, cl_queue, devctx,
                       strides=src.strides, tune=tune)
    if backend == Backend.PYOPENCL:
        app.fft(src, dest, queue=cl_queue)
    else:
        app.fft(src, dest)
    if return_scale:
        s = app.get_fft_scale()
        return dest, s
    return dest


def ifftn(src, dest=None, ndim=None, norm=1, axes=None, cuda_stream=None, cl_queue=None,
          return_scale=False, tune=False):
    """
    Perform an inverse FFT on a GPU array, automatically creating the VkFFTApp
    and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: if 0 (un-normalised), every transform multiplies the L2 norm
        of the array by the transform size.
        if 1 (the default) or "backward", the inverse transform divides the
        L2 norm by the array size, so FFT+iFFT will keep the array norm.
        if "ortho", each transform will keep the L2 norm, but that will
        involve an extra read & write operation.
    :param axes: a list or tuple of axes along which the transform is made.
        if None, the transform is done along the ndim fastest axes, or all
        axes if ndim is None. Not allowed for R2C transforms
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param return_scale: if True, return the scale factor by which the result
        must be multiplied to keep its L2 norm after the transform
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array if return_scale is False, or (dest, scale)
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_fft_app(backend, src.shape, src.dtype, inplace, ndim, axes, norm, cuda_stream, cl_queue, devctx,
                       strides=src.strides, tune=tune)
    if backend == Backend.PYOPENCL:
        app.ifft(src, dest, queue=cl_queue)
    else:
        app.ifft(src, dest)
    if return_scale:
        s = app.get_fft_scale()
        return dest, s
    return dest


def rfftn(src, dest=None, ndim=None, norm=1, cuda_stream=None, cl_queue=None,
          return_scale=False, tune=False, r2c_odd=False):
    """
    Perform a real->complex transform on a GPU array, automatically creating
    the VkFFTApp and caching it for future re-use.
    For an out-of-place transform, the length of the destination last axis will
    be src.shape[-1]//2+1.
    For an in-place transform with an even [respectively odd]-sized
    fast (x) axis, the src array should have a shape (..., nx+2)
    [respectively (..., nx+1)], the last one or two values along the
    fast (x) axis are ignored, and the destination
    array will have a shape of (..., nx//2+1).
    An in-place transform with an odd-sized x-axis requires r2c_odd=True

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: if 0 (un-normalised), every transform multiplies the L2 norm
        of the array by the transform size.
        if 1 (the default) or "backward", the inverse transform divides the
        L2 norm by the array size, so FFT+iFFT will keep the array norm.
        if "ortho", each transform will keep the L2 norm, but that will
        involve an extra read & write operation.
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param return_scale: if True, return the scale factor by which the result
        must be multiplied to keep its L2 norm after the transform
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :param r2c_odd: should be set to True for an in-place r2c transform
        where the actual data length is odd along the fast axis. This
        parameter is ignored otherwise.
    :return: the destination array if return_scale is False, or (dest, scale).
        For an in-place transform, the returned value is a view of the array
        with the appropriate type.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx, dtype = \
        _prepare_transform(src, dest, cl_queue, cuda_stream, True)
    app = _get_rfft_app(backend, src.shape, src.dtype, inplace, ndim, norm, cuda_stream, cl_queue, devctx,
                        strides=src.strides, tune=tune, r2c_odd=r2c_odd)
    if backend == Backend.PYOPENCL:
        app.fft(src, dest, queue=cl_queue)
    else:
        app.fft(src, dest)
    if return_scale:
        s = app.get_fft_scale()
        return dest.view(dtype=dtype), s
    return dest.view(dtype=dtype)


def irfftn(src, dest=None, ndim=None, norm=1, cuda_stream=None, cl_queue=None,
           return_scale=False, tune=False, r2c_odd=False):
    """
    Perform a complex->real transform on a GPU array, automatically creating
    the VkFFTApp and caching it for future re-use.
    For an out-of-place transform, the length of the destination last axis will
    be (src.shape[-1]-1)*2.
    For an in-place transform, if the src complex array has a shape (..., nx),
    the destination (real) array will have a shape of (..., nx*2), but the last
    one (if r2c_odd=True) or two values along the x-axis are used as buffer:
    the size of the transform is thus either nx*2 or nx*2+1.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: if 0 (un-normalised), every transform multiplies the L2 norm
        of the array by the transform size.
        if 1 (the default) or "backward", the inverse transform divides the
        L2 norm by the array size, so FFT+iFFT will keep the array norm.
        if "ortho", each transform will keep the L2 norm, but that will
        involve an extra read & write operation.
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param return_scale: if True, return the scale factor by which the result
        must be multiplied to keep its L2 norm after the transform
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :param r2c_odd: should be set to True for an in-place r2c transform
        where the actual data length (in the real array) is odd along the
        fast axis. This parameter is ignored otherwise.
    :return: the destination array if return_scale is False, or (dest, scale)
        For an in-place transform, the returned value is a view of the array
        with the appropriate type.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx, dtype = \
        _prepare_transform(src, dest, cl_queue, cuda_stream, True, r2c_odd=r2c_odd)
    app = _get_rfft_app(backend, dest.shape, dest.dtype, inplace, ndim, norm, cuda_stream, cl_queue, devctx,
                        strides=dest.strides, tune=tune, r2c_odd=r2c_odd)
    if backend == Backend.PYOPENCL:
        app.ifft(src, dest, queue=cl_queue)
    else:
        app.ifft(src, dest)
    if return_scale:
        s = app.get_fft_scale()
        return dest.view(dtype=dtype), s
    return dest.view(dtype=dtype)


def dctn(src, dest=None, ndim=None, norm=1, dct_type=2, cuda_stream=None, cl_queue=None, tune=False):
    """
    Perform a real->real Direct Cosine Transform on a GPU array, automatically
    creating the VkFFTApp and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: normalisation mode, either 0 (un-normalised) or
        1 (the default, also available as "backward) which will normalise
        the inverse transform, so DCT+iDCT will keep the array norm.
    :param dct_type: the type of dct desired: 1, 2 (default), 3 or 4
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_dct_app(backend, src.shape, src.dtype, inplace, ndim, norm,
                       dct_type, cuda_stream, cl_queue, devctx, tune=tune)
    if backend == Backend.PYOPENCL:
        app.fft(src, dest, queue=cl_queue)
    else:
        app.fft(src, dest)
    return dest


def idctn(src, dest=None, ndim=None, norm=1, dct_type=2, cuda_stream=None, cl_queue=None, tune=False):
    """
    Perform a real->real inverse Direct Cosine Transform on a GPU array,
    automatically creating the VkFFTApp and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: normalisation mode, either 0 (un-normalised) or
        1 (the default, also available as "backward) which will normalise
        the inverse transform, so DCT+iDCT will keep the array norm.
    :param dct_type: the type of dct desired: 2 (default), 3 or 4
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_dct_app(backend, src.shape, src.dtype, inplace, ndim, norm,
                       dct_type, cuda_stream, cl_queue, devctx, tune=tune)
    if backend == Backend.PYOPENCL:
        app.ifft(src, dest, queue=cl_queue)
    else:
        app.ifft(src, dest)
    return dest


def dstn(src, dest=None, ndim=None, norm=1, dst_type=2, cuda_stream=None, cl_queue=None, tune=False):
    """
    Perform a real->real Direct Cosine Transform on a GPU array, automatically
    creating the VkFFTApp and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: normalisation mode, either 0 (un-normalised) or
        1 (the default, also available as "backward) which will normalise
        the inverse transform, so DST+iDST will keep the array norm.
    :param dst_type: the type of dst desired: 1, 2 (default), 3 or 4
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_dst_app(backend, src.shape, src.dtype, inplace, ndim, norm,
                       dst_type, cuda_stream, cl_queue, devctx, tune=tune)
    if backend == Backend.PYOPENCL:
        app.fft(src, dest, queue=cl_queue)
    else:
        app.fft(src, dest)
    return dest


def idstn(src, dest=None, ndim=None, norm=1, dst_type=2, cuda_stream=None, cl_queue=None, tune=False):
    """
    Perform a real->real inverse Direct Cosine Transform on a GPU array,
    automatically creating the VkFFTApp and caching it for future re-use.

    :param src: the source pycuda.gpuarray.GPUArray or cupy.ndarray
    :param dest: the destination GPU array. If None, a new GPU array will
        be created and returned (using the source array allocator
        (pycuda, pyopencl) if available).
        If dest is the same array as src, an inplace transform is done.
    :param ndim: the number of dimensions (<=3) to use for the FFT. By default,
        uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D
        array to perform a batched 3D FFT on all the layers. The FFT
        is always performed along the last axes if the array's number
        of dimension is larger than ndim, i.e. on the x-axis for ndim=1,
        on the x and y axes for ndim=2.
    :param norm: normalisation mode, either 0 (un-normalised) or
        1 (the default, also available as "backward) which will normalise
        the inverse transform, so DST+iDST will keep the array norm.
    :param dst_type: the type of dst desired: 2 (default), 3 or 4
    :param cuda_stream: the pycuda.driver.Stream or cupy.cuda.Stream to use
        for the transform. If None, the default one will be used
    :param cl_queue: the pyopencl.CommandQueue to be used. If None,
        the source array default queue will be used
    :param tune: if True, will activate the automatic tuning of VkFFT
        parameters to maximise the FT throughput. This uses a quick
        approach testing a few transforms (about 4) before choosing the
        optimal parameters. This is similar to FFTW's FFTW_MEASURE approach.
    :return: the destination array.
    """
    backend, inplace, dest, cl_queue, cuda_stream, devctx = _prepare_transform(src, dest, cl_queue, cuda_stream, False)
    app = _get_dst_app(backend, src.shape, src.dtype, inplace, ndim, norm,
                       dst_type, cuda_stream, cl_queue, devctx, tune=tune)
    if backend == Backend.PYOPENCL:
        app.ifft(src, dest, queue=cl_queue)
    else:
        app.ifft(src, dest)
    return dest


def clear_vkfftapp_cache():
    """ Remove all cached VkFFTApp"""
    _get_fft_app.cache_clear()
    _get_rfft_app.cache_clear()
