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

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

import ctypes
import os
import platform
import sysconfig
import warnings
from enum import Enum

import numpy as np

import radis.gpu.vulkan.pyvkfft_config as config

# import pyvkfft_config as config


# np.complex32 does not exist yet https://github.com/numpy/numpy/issues/14753
complex32 = np.dtype([("re", np.float16), ("im", np.float16)])

# Type to pass array size, omit and batch as int arrays
ctype_int_size_p = np.ctypeslib.ndpointer(dtype=int, ndim=1, flags="C_CONTIGUOUS")


class VkFFTResult(Enum):
    """VkFFT error codes from vkFFT.h"""

    VKFFT_SUCCESS = 0
    VKFFT_ERROR_MALLOC_FAILED = 1
    VKFFT_ERROR_INSUFFICIENT_CODE_BUFFER = 2
    VKFFT_ERROR_INSUFFICIENT_TEMP_BUFFER = 3
    VKFFT_ERROR_PLAN_NOT_INITIALIZED = 4
    VKFFT_ERROR_NULL_TEMP_PASSED = 5
    VKFFT_ERROR_INVALID_PHYSICAL_DEVICE = 1001
    VKFFT_ERROR_INVALID_DEVICE = 1002
    VKFFT_ERROR_INVALID_QUEUE = 1003
    VKFFT_ERROR_INVALID_COMMAND_POOL = 1004
    VKFFT_ERROR_INVALID_FENCE = 1005
    VKFFT_ERROR_ONLY_FORWARD_FFT_INITIALIZED = 1006
    VKFFT_ERROR_ONLY_INVERSE_FFT_INITIALIZED = 1007
    VKFFT_ERROR_INVALID_CONTEXT = 1008
    VKFFT_ERROR_INVALID_PLATFORM = 1009
    VKFFT_ERROR_ENABLED_saveApplicationToString = (1010,)
    VKFFT_ERROR_EMPTY_FFTdim = 2001
    VKFFT_ERROR_EMPTY_size = 2002
    VKFFT_ERROR_EMPTY_bufferSize = 2003
    VKFFT_ERROR_EMPTY_buffer = 2004
    VKFFT_ERROR_EMPTY_tempBufferSize = 2005
    VKFFT_ERROR_EMPTY_tempBuffer = 2006
    VKFFT_ERROR_EMPTY_inputBufferSize = 2007
    VKFFT_ERROR_EMPTY_inputBuffer = 2008
    VKFFT_ERROR_EMPTY_outputBufferSize = 2009
    VKFFT_ERROR_EMPTY_outputBuffer = 2010
    VKFFT_ERROR_EMPTY_kernelSize = 2011
    VKFFT_ERROR_EMPTY_kernel = 2012
    VKFFT_ERROR_EMPTY_applicationString = (2013,)
    VKFFT_ERROR_UNSUPPORTED_RADIX = 3001
    VKFFT_ERROR_UNSUPPORTED_FFT_LENGTH = 3002
    VKFFT_ERROR_UNSUPPORTED_FFT_LENGTH_R2C = 3003
    VKFFT_ERROR_UNSUPPORTED_FFT_LENGTH_DCT = 3004
    VKFFT_ERROR_UNSUPPORTED_FFT_OMIT = 3005
    VKFFT_ERROR_FAILED_TO_ALLOCATE = 4001
    VKFFT_ERROR_FAILED_TO_MAP_MEMORY = 4002
    VKFFT_ERROR_FAILED_TO_ALLOCATE_COMMAND_BUFFERS = 4003
    VKFFT_ERROR_FAILED_TO_BEGIN_COMMAND_BUFFER = 4004
    VKFFT_ERROR_FAILED_TO_END_COMMAND_BUFFER = 4005
    VKFFT_ERROR_FAILED_TO_SUBMIT_QUEUE = 4006
    VKFFT_ERROR_FAILED_TO_WAIT_FOR_FENCES = 4007
    VKFFT_ERROR_FAILED_TO_RESET_FENCES = 4008
    VKFFT_ERROR_FAILED_TO_CREATE_DESCRIPTOR_POOL = 4009
    VKFFT_ERROR_FAILED_TO_CREATE_DESCRIPTOR_SET_LAYOUT = 4010
    VKFFT_ERROR_FAILED_TO_ALLOCATE_DESCRIPTOR_SETS = 4011
    VKFFT_ERROR_FAILED_TO_CREATE_PIPELINE_LAYOUT = 4012
    VKFFT_ERROR_FAILED_SHADER_PREPROCESS = 4013
    VKFFT_ERROR_FAILED_SHADER_PARSE = 4014
    VKFFT_ERROR_FAILED_SHADER_LINK = 4015
    VKFFT_ERROR_FAILED_SPIRV_GENERATE = 4016
    VKFFT_ERROR_FAILED_TO_CREATE_SHADER_MODULE = 4017
    VKFFT_ERROR_FAILED_TO_CREATE_INSTANCE = 4018
    VKFFT_ERROR_FAILED_TO_SETUP_DEBUG_MESSENGER = 4019
    VKFFT_ERROR_FAILED_TO_FIND_PHYSICAL_DEVICE = 4020
    VKFFT_ERROR_FAILED_TO_CREATE_DEVICE = 4021
    VKFFT_ERROR_FAILED_TO_CREATE_FENCE = 4022
    VKFFT_ERROR_FAILED_TO_CREATE_COMMAND_POOL = 4023
    VKFFT_ERROR_FAILED_TO_CREATE_BUFFER = 4024
    VKFFT_ERROR_FAILED_TO_ALLOCATE_MEMORY = 4025
    VKFFT_ERROR_FAILED_TO_BIND_BUFFER_MEMORY = 4026
    VKFFT_ERROR_FAILED_TO_FIND_MEMORY = 4027
    VKFFT_ERROR_FAILED_TO_SYNCHRONIZE = 4028
    VKFFT_ERROR_FAILED_TO_COPY = 4029
    VKFFT_ERROR_FAILED_TO_CREATE_PROGRAM = 4030
    VKFFT_ERROR_FAILED_TO_COMPILE_PROGRAM = 4031
    VKFFT_ERROR_FAILED_TO_GET_CODE_SIZE = 4032
    VKFFT_ERROR_FAILED_TO_GET_CODE = 4033
    VKFFT_ERROR_FAILED_TO_DESTROY_PROGRAM = 4034
    VKFFT_ERROR_FAILED_TO_LOAD_MODULE = 4035
    VKFFT_ERROR_FAILED_TO_GET_FUNCTION = 4036
    VKFFT_ERROR_FAILED_TO_SET_DYNAMIC_SHARED_MEMORY = 4037
    VKFFT_ERROR_FAILED_TO_MODULE_GET_GLOBAL = 4038
    VKFFT_ERROR_FAILED_TO_LAUNCH_KERNEL = 4039
    VKFFT_ERROR_FAILED_TO_EVENT_RECORD = 4040
    VKFFT_ERROR_FAILED_TO_ADD_NAME_EXPRESSION = 4041
    VKFFT_ERROR_FAILED_TO_INITIALIZE = 4042
    VKFFT_ERROR_FAILED_TO_SET_DEVICE_ID = 4043
    VKFFT_ERROR_FAILED_TO_GET_DEVICE = 4044
    VKFFT_ERROR_FAILED_TO_CREATE_CONTEXT = 4045
    VKFFT_ERROR_FAILED_TO_CREATE_PIPELINE = 4046
    VKFFT_ERROR_FAILED_TO_SET_KERNEL_ARG = 4047
    VKFFT_ERROR_FAILED_TO_CREATE_COMMAND_QUEUE = 4048
    VKFFT_ERROR_FAILED_TO_RELEASE_COMMAND_QUEUE = 4049
    VKFFT_ERROR_FAILED_TO_ENUMERATE_DEVICES = 4050
    VKFFT_ERROR_FAILED_TO_GET_ATTRIBUTE = 4051
    VKFFT_ERROR_FAILED_TO_CREATE_EVENT = 4052


def load_library(basename):
    if platform.system() == "Windows":
        # We patched build_ext so the module is a .so and not a dll
        ext = ".so"
    else:
        ext = sysconfig.get_config_var("EXT_SUFFIX")
    return ctypes.cdll.LoadLibrary(
        os.path.join(os.path.dirname(__file__) or os.path.curdir, basename + ext)
    )


def primes(n):
    """Returns the prime decomposition of n as a list.
    This only remains as a useful function, but VkFFT
    allows any prime decomposition, even if performance
    can be better for prime numbers <=13.
    """
    v = [1]
    assert n > 0
    i = 2
    while i * i <= n:
        while n % i == 0:
            v.append(i)
            n //= i
        i += 1
    if n > 1:
        v.append(n)
    return v


def radix_gen(
    nmax, radix, even=False, exclude_one=True, inverted=False, nmin=None, max_pow=None
):
    """
    Generate an array of integers which are only multiple of powers
    of base integers, e.g. 2**N1 * 3**N2 * 5**N3 etc...

    :param nmax: the maximum integer to return (included)
    :param radix: the list/tuple of base integers - which don't need
        to be primes
    :param even: if True, only return even numbers
    :param exclude_one: if True (the default), exclude 1
    :param inverted: if True, the returned array will only include
        integers which are NOT in the form 2**N1 * 3**N2 * 5**N3...
    :param nmin: if not None, the integer values returned will be >=nmin
    :param max_pow: if not None, the N1, N2, N3... powers (for sizes
        in the form 2**N1 * 3**N2 * 5**N3) will be at max equal to this
        value, which allows to reduce the number of generated sizes
        while testing all base radixes
    :return: the numpy array of integers, sorted
    """
    a = np.ones(1, dtype=np.int64)
    for i in range(len(radix)):
        if max_pow is None:
            tmp = np.arange(int(np.floor(np.log(nmax) / np.log(radix[i]))) + 1)
        else:
            tmp = np.arange(
                min(max_pow, int(np.floor(np.log(nmax) / np.log(radix[i])))) + 1
            )
        a = a * radix[i] ** tmp[:, np.newaxis]
        a = a.flatten()
        a = a[a <= nmax]
    if inverted:
        b = np.arange(nmax + 1)
        b[a] = 0
        a = b.take(np.nonzero(b)[0])
    if even:
        a = a[(a % 2) == 0]
    if nmin is not None:
        a = a[a >= nmin]
    a.sort()
    if len(a):
        if exclude_one and a[0] == 1:
            return a[1:]
    return a


def radix_gen_n(
    nmax,
    max_size,
    radix,
    ndim=None,
    even=False,
    exclude_one=True,
    inverted=False,
    nmin=None,
    max_pow=None,
    range_nd_narrow=None,
    min_size=0,
):
    """
    Generate a list of array shape with integers which are only multiple
    of powers of base integers, e.g. 2**N1 * 3**N2 * 5**N3 etc...,
    for each of the dimensions, and with a maximum size.
    Note that this can generate a large number of sizes.

    :param nmax: the maximum value for the length of each dimension (included)
    :param max_size: the maximum size (number of elements) for the array
    :param radix: the list/tuple of base integers - which don't need
        to be primes. If None, all sizes are allowed
    :param ndim: the number of dimensions allowed. If None, 1D, 2D and 3D
        shapes are mixed.
    :param even: if True, only return even numbers
    :param exclude_one: if True (the default), exclude 1
    :param inverted: if True, the returned array will only include
        integers which are NOT in the form 2**N1 * 3**N2 * 5**N3...
    :param nmin: if not None, the integer values returned will be >=nmin
    :param max_pow: if not None, the N1, N2, N3... powers (for sizes
        in the form 2**N1 * 3**N2 * 5**N3) will be at max equal to this
        value, which allows to reduce the number of generated sizes
        while testing all base radixes
    :param range_nd_narrow: if a tuple of values (drel, dabs) is given,
        with drel within [0;1], for dimensions>1,
        in an array of shape (s0, s1, s2), the difference of lengths
        with respect to the first dimension cannot be larger than
        min(drel * s0, dabs). This allows to reduce the number
        of shapes tested. With drel=dabs=0, all dimensions must
        have identical lengths.
    :param min_size: the minimum size (number of elements). This can be
        used to separate large array tests and use a larger number of
        parallel process for smaller ones.
    :return: the list of array shapes.
    """
    v = []
    if radix is None:
        if even:
            if nmin is None:
                n0 = np.arange(2, min(nmax, max_size), 2)
            else:
                n0 = np.arange(nmin + nmin % 2, min(nmax, max_size), 2)
        else:
            if nmin is None:
                n0 = np.arange(2, min(nmax, max_size))
            else:
                n0 = np.arange(nmin, min(nmax, max_size))
    else:
        n0 = radix_gen(
            nmax,
            radix,
            even=even,
            exclude_one=exclude_one,
            inverted=inverted,
            nmin=nmin,
            max_pow=max_pow,
        )

    if ndim is None or ndim in [1, 12, 123]:
        idx = np.nonzero((n0 <= max_size) * (n0 >= min_size))[0]
        if len(idx):
            v += list(zip(n0.take(idx)))
    if ndim is None or ndim in [2, 12, 123]:
        vidx = list(range(0, len(n0), 1000))
        if vidx[-1] != len(n0):
            vidx.append(len(n0))
        for i1 in range(len(vidx) - 1):
            l01 = n0[vidx[i1] : vidx[i1 + 1]]
            for i2 in range(len(vidx) - 1):
                l2 = n0[vidx[i2] : vidx[i2 + 1]][:, np.newaxis]
                s = (l01 * l2).flatten()
                l1, l2 = (l01 + np.zeros_like(l2)).flatten(), (
                    l2 + np.zeros_like(l01)
                ).flatten()
                tmp = (s <= max_size) * (s >= min_size)
                if range_nd_narrow is not None:
                    drel, dabs = range_nd_narrow
                    m = np.maximum(dabs, l1 * drel)
                    tmp = np.logical_and(tmp, abs(l1 - l2) <= m)
                idx = np.nonzero(tmp)[0]
                if len(idx):
                    v += list(zip(l1.take(idx), l2.take(idx)))
    if ndim is None or ndim in [3, 123]:
        vidx = list(range(0, len(n0), 100))
        if vidx[-1] != len(n0):
            vidx.append(len(n0))
        for i1 in range(len(vidx) - 1):
            l01 = n0[vidx[i1] : vidx[i1 + 1]]
            for i2 in range(len(vidx) - 1):
                l02 = n0[vidx[i2] : vidx[i2 + 1], np.newaxis]
                for i3 in range(len(vidx) - 1):
                    l3 = n0[vidx[i3] : vidx[i3 + 1], np.newaxis, np.newaxis]
                    # print(i1, i2, i3, l1.shape, l2.shape, l3.shape)
                    s = (l01 * l02 * l3).flatten()
                    l1, l2, l3 = (
                        (l01 + np.zeros_like(l02) + np.zeros_like(l3)).flatten(),
                        (l02 + np.zeros_like(l01) + np.zeros_like(l3)).flatten(),
                        (l3 + np.zeros_like(l01) + np.zeros_like(l02)).flatten(),
                    )
                    tmp = (s <= max_size) * (s >= min_size)
                    if range_nd_narrow is not None:
                        drel, dabs = range_nd_narrow
                        m = np.maximum(dabs, l1 * drel)
                        tmp = np.logical_and(
                            tmp, (abs(l1 - l2) <= m) * (abs(l1 - l3) <= m)
                        )
                    idx = np.nonzero(tmp)[0]
                    if len(idx):
                        v += list(zip(l1.take(idx), l2.take(idx), l3.take(idx)))
    return v


def calc_transform_axes(shape, axes=None, ndim=None, strides=None):
    """Compute the final shape of the array to be passed
    to VkFFT, and the axes for which the transform should
    be skipped.
    By collapsing non-transformed consecutive axes and using batch transforms,
    it is possible to support larger dimensions (the limit is set at
    compilation time).

    :param shape: the initial shape of the data array. Note that this shape
        should be in the usual numpy order, i.e. the fastest axis is
        listed last. e.g. (nz, ny, nx)
    :param axes: the axes to be transformed. if None, all axes
        are transformed, or up to ndim.
    :param ndim: the number of dimensions for the transform. If None,
        the number of axes is used
    :param strides: the array strides. If None, a C-order is assumed
        with the fastest axes along the last dimensions (numpy default)
    :return: (shape, n_batch, skip_axis, ndim) with the shape after collapsing
        consecutive non-transformed axes (padded with ones if necessary,
        with the order adequate for VkFFT i.e. (nx, ny, nz,...),
        the batch size (e.g. 5 if shape=(5,16,16) and ndim=2),
        the list of booleans indicating which axes should
        be skipped, and the number of transform axes.
    """
    # reverse order to have list as (nx, ny, nz,...)
    shape1 = list(reversed(list(shape)))
    if np.isscalar(axes):
        axes = [axes]
    if ndim is None:
        if axes is None:
            ndim1 = len(shape)
            axes = list(range(ndim1))
    else:
        ndim1 = ndim
        if axes is None:
            axes = list(range(-1, -ndim - 1, -1))
        else:
            if ndim1 != len(axes):
                raise RuntimeError(
                    "The number of transform axes does not match ndim:", axes, ndim
                )

    if strides is not None and len(shape) > 1:
        idx = np.argsort(strides)[::-1]
        if not np.alltrue((idx[1:] - idx[:-1]) > 0):
            # Array is not C-ordered, need to move axes
            idxend = -len(idx) + idx
            # print("Re-ordering:", shape1, np.take(shape1, idx), "axes:", axes, [idxend[i] for i in axes])
            shape1 = np.take(shape1, idx).tolist()
            axes = [idxend[i] for i in axes]

    # Collapse non-transform axes when possible
    skip_axis = [True] * len(shape1)
    for i in axes:
        skip_axis[i] = False
    skip_axis = list(reversed(skip_axis))  # numpy (z,y,x) to vkfft (x,y,z) order
    # print(shape1, skip_axis, axes)
    i = 0
    while i <= len(shape1) - 2:
        if skip_axis[i] and skip_axis[i + 1]:
            shape1[i] *= shape1[i + 1]
            shape1.pop(i + 1)
            skip_axis.pop(i + 1)
        else:
            i += 1

    # Fix ndim so skipped axes are counted
    ndim1 = len(shape1) - list(reversed(skip_axis)).index(False)

    # Axes beyond ndim are marked skipped
    for i in range(ndim1, len(shape1)):
        skip_axis[i] = False

    # For VkFFT > cc2b427, all dimensions beyond the
    # transformed axes should be in n_batch
    n_batch = 1
    if len(shape1) > ndim1:
        for i in range(ndim1, len(shape1)):
            n_batch *= shape1[i]
            shape1[i] = 1
    # print(shape, axes, ndim, strides, "->", shape1, skip_axis)
    return shape1, n_batch, skip_axis, ndim1


def check_vkfft_result(
    res,
    shape=None,
    dtype=None,
    ndim=None,
    inplace=None,
    norm=None,
    r2c=None,
    dct=None,
    axes=None,
    backend=None,
):
    """
    Check VkFFTResult code.

    :param res: the result code from launching a transform.
    :param shape: shape of the array
    :param dtype: data type of the array
    :param ndim: number of transform dimensions
    :param inplace: True or False
    :param norm: 0 or1 or "ortho"
    :param r2c: True or False
    :param dct: False, 1, 2, 3 or 4
    :param axes: transform axes
    :param backend: the backend
    :raises RuntimeError: if res != 0
    """
    if isinstance(res, ctypes.c_int):
        res = res.value
    if res != 0:
        s = ""
        if r2c:
            s += "R2C "
        elif dct:
            s += "DCT%d " % dct
        else:
            s += "C2C "
        if r2c and inplace and shape is not None:
            tmp = list(shape)
            tmp[-1] -= 2
            shstr = str(tuple(tmp)).replace(" ", "")
            if ",)" in shstr:
                s += shstr.replace(",)", "+2)") + " "
            else:
                s += shstr.replace(")", "+2)") + " "
        else:
            s += str(shape).replace(" ", "") + " "
        if dtype is not None:
            s += str(dtype) + " "
        if axes is not None:
            s += str(axes).replace(" ", "") + " "
        if ndim is not None:
            s += "%dD " % ndim
        if inplace:
            s += "inplace "
        if norm:
            s += "norm=%s " % str(norm)
        if backend is not None:
            s += "[%s]" % backend
        try:
            r = VkFFTResult(res)
            raise RuntimeError("VkFFT error %d: %s %s" % (res, r.name, s))
        except ValueError:
            raise RuntimeError("VkFFT error %d (unknown) %s" % (res, s))


class VkFFTApp:
    """
    VkFFT application interface implementing a FFT plan
    """

    def __init__(
        self,
        shape,
        dtype: type,
        ndim=None,
        inplace=True,
        norm=1,
        r2c=False,
        dct=False,
        axes=None,
        strides=None,
        **kwargs
    ):
        """
        Init function for the VkFFT application.
        :raises RuntimeError:  if the transform dimensions are not allowed by VkFFT.
        """
        self.app = None
        self.config = None
        if dct and r2c:
            raise RuntimeError("R2C and DCT cannot both be selected !")
        if (r2c or dct) and dtype not in [np.float16, np.float32, np.float64]:
            raise RuntimeError("R2C or DCT selected but input type is not real !")
        if r2c and axes is not None:
            raise RuntimeError("axes=... is not allowed for R2C transforms")
        if strides is not None and dct and len(shape) > 1:
            # TODO: update support status for non-C-contiguous DCT transforms
            s = np.array(strides)
            if not np.alltrue((s[:-1] - s[1:]) > 0):
                raise RuntimeError(
                    "A C-contiguous array is required for DCT transforms"
                )
        # Get the final shape passed to VkFFT, collapsing non-transform axes
        # as necessary. The calculated shape has 4 dimensions (nx, ny, nz, n_batch).
        # self.shape, self.n_batch, self.skip_axis, self.ndim = calc_transform_axes(shape, axes, ndim, strides)
        self.shape = shape
        self.ndim = ndim

        self.inplace = inplace
        self.r2c = r2c
        if dct is False:
            self.dct = 0
        elif dct is True:
            self.dct = 2
        else:
            self.dct = dct
        if dct and self.dct < 1 or dct > 4:
            raise RuntimeError("Only DCT of types 1, 2, 3 and 4 are allowed")
        # print("VkFFTApp:", shape, axes, ndim, "->", self.shape, self.skip_axis, self.ndim)

        # Experimental parameters. Not much difference is seen, so don't document this,
        # VkFFT default parameters seem fine.
        if "disableReorderFourStep" in kwargs:
            self.disableReorderFourStep = kwargs["disableReorderFourStep"]
        else:
            self.disableReorderFourStep = -1

        # uint64_t coalescedMemory - number of bytes to coalesce per one transaction.
        # For Nvidia and AMD is equal to 32, Intel is equal to 64. Going to work regardless,
        # but if specified by the user correctly, the performance will be higher. Default 64
        # for other GPUs. For half-precision should be multiplied by two. Should be a power of two.
        if "coalescedMemory" in kwargs:
            self.coalescedMemory = kwargs["coalescedMemory"]
        else:
            self.coalescedMemory = -1

        # uint64_t numSharedBanks - configure the number of shared banks on the target GPU.
        # Default 32. Minor performance boost as it solves shared memory conflicts for
        # the power of two systems.
        if "numSharedBanks" in kwargs:
            self.numSharedBanks = kwargs["numSharedBanks"]
        else:
            self.numSharedBanks = -1

        # uint64_t aimThreads - try to aim all kernels at this amount of threads.
        # Gains/losses are not predictable, just a parameter to play with
        # (it is not guaranteed that the target kernel will use that many threads).
        # Default 128.
        if "aimThreads" in kwargs:
            self.aimThreads = kwargs["aimThreads"]
        else:
            self.aimThreads = -1

        # uint64_t performBandwidthBoost - try to reduce coalesced number by a factor of X
        # to get bigger sequence in one upload for strided axes.
        # Default: -1(inf) for DCT, 2 for Bluestein’s algorithm (or -1 if DCT), 0 otherwise
        if "performBandwidthBoost" in kwargs:
            self.performBandwidthBoost = kwargs["performBandwidthBoost"]
        else:
            self.performBandwidthBoost = -1

        # uint64_t registerBoost - specify if the register file size is bigger than
        # shared memory and can be used to extend it X times (on Nvidia 256KB register
        # file can be used instead of 32KB of shared memory, set this constant to 4 to
        # emulate 128KB of shared memory). Default 1 - no over-utilization. In Vulkan,
        # OpenCL and Level Zero it is set to 4 on Nvidia GPUs, to 2 if the driver
        # shows 64KB or more of shared memory on AMD, to 2 if the driver shows less
        # than 64KB of shared memory on AMD, to 1 if the driver shows 64KB or more
        # of shared memory on Intel, to 2 if the driver shows less than 64KB of shared
        # memory on Intel.
        if "registerBoost" in kwargs:
            self.registerBoost = kwargs["registerBoost"]
        else:
            self.registerBoost = -1

        # uint64_t registerBoostNonPow2 - specify if register over-utilization should
        # be used on non-power of 2 sequences. Default 0, set to 1 to enable.
        if "registerBoostNonPow2" in kwargs:
            self.registerBoostNonPow2 = kwargs["registerBoostNonPow2"]
        else:
            self.registerBoostNonPow2 = -1

        # uint64_t registerBoost4Step - specify if register file over-utilization
        # should be used in big sequences (>2^14), same definition as registerBoost.
        # Default 1.
        if "registerBoost4Step" in kwargs:
            self.registerBoost4Step = kwargs["registerBoost4Step"]
        else:
            self.registerBoost4Step = -1

        if "warpSize" in kwargs:
            self.warpSize = kwargs["warpSize"]
        else:
            self.warpSize = -1

        self.groupedBatch = [-1, -1, -1]
        if "groupedBatch" in kwargs:
            self.groupedBatch = list(kwargs["groupedBatch"])
            # In case less than 3 parameters where given
            self.groupedBatch[: len(kwargs["groupedBatch"])] = kwargs["groupedBatch"]

        if "useLUT" in kwargs:
            # useLUT=1 may be beneficial on platforms which have a low accuracy for
            # the native sincos functions.
            if kwargs["useLUT"] is None:
                self.use_lut = -1
            else:
                self.use_lut = kwargs["useLUT"]
        elif config.USE_LUT is not None:
            self.use_lut = config.USE_LUT
        else:
            self.use_lut = -1

        if "keepShaderCode" in kwargs:
            # This will print the compiled code if equal to 1
            self.keepShaderCode = kwargs["keepShaderCode"]
        else:
            self.keepShaderCode = -1

        if norm == "backward":
            norm = 1
        self.norm = norm

        # Precision: number of bytes per float
        if dtype in [np.float16, complex32]:
            self.precision = 2
        elif dtype in [np.float32, np.complex64]:
            self.precision = 4
        elif dtype in [np.float64, np.complex128]:
            self.precision = 8

    def _get_fft_scale(self, norm):
        """Return the scale factor by which an array must be multiplied to keep its L2 norm
        after a forward FT
        :param norm: the norm option for which the scale is computed, either 0 or 1
        :return: the scale factor, as a numpy float with the precision used for the fft
        """
        dtype = np.float32
        if self.precision == 8:
            dtype = np.float64
        elif self.precision == 2:
            dtype = np.float16
        s = 1
        ndim_real = 0
        for i in range(self.ndim):
            if not self.skip_axis[i]:
                s *= self.shape[i]
                ndim_real += 1
        s = np.sqrt(s)
        if self.r2c and self.inplace:
            # Note: this is still correct for non-C-ordered arrays since
            # self.shape has been re-ordered
            s *= np.sqrt((self.shape[0] - 2) / self.shape[0])
        if self.dct:
            s *= 2 ** (0.5 * ndim_real)
            if self.dct != 4:
                warnings.warn(
                    "A DCT type 2 or 3 cannot be strictly normalised, using approximation,"
                    " see https://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-II"
                )
        if norm == 0 or norm == 1:
            return dtype(1 / s)
        elif norm == "ortho":
            return dtype(1)
        raise RuntimeError("Unknown norm choice !")

    def get_fft_scale(self):
        """Return the scale factor by which an array must be multiplied to keep its L2 norm
        after a forward FT
        """
        return self._get_fft_scale(self.norm)

    def _get_ifft_scale(self, norm):
        """Return the scale factor by which an array must be multiplied to keep its L2 norm
        after a backward FT
        :param norm: the norm option for which the scale is computed, either 0 or 1
        :return: the scale factor, as a numpy float with the precision used for the fft
        """
        dtype = np.float32
        if self.precision == 8:
            dtype = np.float64
        elif self.precision == 2:
            dtype = np.float16
        s = 1
        s_dct = 1
        for i in range(self.ndim):
            if not self.skip_axis[i]:
                s *= self.shape[i]
                if self.dct:
                    s_dct *= np.sqrt(2)
        s = np.sqrt(s)
        if self.r2c and self.inplace:
            s *= np.sqrt((self.shape[0] - 2) / self.shape[0])
        if self.dct and self.dct != 4:
            warnings.warn(
                "A DCT type 2 or 3 cannot be strictly normalised, using approximation,"
                " see https://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-II"
            )
        if norm == 0:
            return dtype(1 / (s * s_dct))
        elif norm == 1:
            # Not sure why the difference in scale factors
            if self.dct == 2:
                s_dct = s_dct**1
            elif self.dct == 3:
                s_dct = s_dct**2
            elif self.dct == 4:
                s_dct = s_dct**3
            return dtype(s * s_dct)
        elif norm == "ortho":
            return dtype(1)
        raise RuntimeError("Unknown norm choice !")

    def get_ifft_scale(self):
        """Return the scale factor by which an array must be multiplied to keep its L2 norm
        after a backward FT
        """
        return self._get_ifft_scale(self.norm)
