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

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

import os
import platform
import multiprocessing
import timeit
import atexit

import psutil
import numpy as np
from numpy.fft import fftn, ifftn, rfftn, irfftn
from pyvkfft.base import primes, strides_nonzero

try:
    # We prefer scipy over numpy for fft, and we can also test dct
    from scipy.fft import dctn, idctn, fftn, ifftn, rfftn, irfftn, dstn, idstn

    has_dct_ref = True
    has_scipy = True
except ImportError:
    has_dct_ref = False
    has_scipy = False

# pyfftw speed is not good compared to scipy, when using every transform once
# try:
#     from pyfftw.interfaces import scipy_fft as pyfftw_fft
#     if not has_scipy:
#         from pyfftw.interfaces.scipy_fft import dctn, idctn, fftn, ifftn, rfftn, irfftn
#
#     has_pyfftw = True
#     has_dct_ref = True
# except ImportError:
#     has_pyfftw = False

try:
    import pyopencl as cl
    import pyopencl.array as cla
    from pyvkfft.opencl import VkFFTApp as clVkFFTApp

    has_opencl = True
except ImportError:
    has_opencl = False

try:
    from pyvkfft.cuda import VkFFTApp as cuVkFFTApp, has_pycuda, has_cupy

    if has_pycuda:
        import pycuda.driver as cu_drv
        import pycuda.gpuarray as cua

    if has_cupy:
        import cupy as cp
except ImportError:
    has_cupy = False
    has_pycuda = False

# Dictionary of cuda/opencl (device, context). Will be initialised on-demand.
# This is needed for multiprocessing.
# The pyopencl entry is a tuple with (device, context, queue, has_cl_fp64)
gpu_ctx_dic = {}


def init_ctx(backend, gpu_name=None, opencl_platform=None, verbose=False):
    if backend in gpu_ctx_dic:
        return
    if backend == "pycuda":
        if not has_pycuda:
            raise RuntimeError("init_ctx: backend=%s is not available" % backend)
        cu_drv.init()
        d = None
        if gpu_name is not None:
            for i in range(cu_drv.Device.count()):
                if gpu_name.lower() in cu_drv.Device(i).name().lower():
                    d = cu_drv.Device(i)
                    break
        else:
            d = cu_drv.Device(0)
        if d is None:
            if gpu_name is not None:
                raise RuntimeError("Selected backend is pycuda, but no device found (name=%s)" % gpu_name)
            else:
                raise RuntimeError("Selected backend is pycuda, but no device found")
        gpu_ctx_dic["pycuda"] = (d, d.retain_primary_context())
        gpu_ctx_dic["pycuda"][1].push()
        if verbose:
            print("Selected device for pycuda: %s" % d.name())
    elif backend == "pyopencl":
        if not has_opencl:
            raise RuntimeError("init_ctx: backend=%s is not available" % backend)
        d = None
        for p in cl.get_platforms():
            if d is not None:
                break
            if opencl_platform is not None:
                if opencl_platform.lower() not in p.name.lower():
                    continue
            for d0 in p.get_devices():
                if d0.type & cl.device_type.GPU:
                    if gpu_name is not None:
                        if gpu_name.lower() in d0.name.lower():
                            d = d0
                    else:
                        d = d0
                if d is not None:
                    break
        if d is None:
            raise RuntimeError("Selected backend is pyopencl, but no device found (name=%s, platform=%s)" %
                               (gpu_name, opencl_platform))
        cl_ctx = cl.Context([d])
        cq = cl.CommandQueue(cl_ctx)
        gpu_ctx_dic["pyopencl"] = d, cl_ctx, cq, 'cl_khr_fp64' in cq.device.extensions
        if verbose:
            print("Selected device for pyopencl: %s [%s]" % (d.name, p.name))
    elif backend == "cupy":
        if not has_cupy:
            raise RuntimeError("init_ctx: backend=%s is not available" % backend)
        if gpu_name is not None:
            for i in range(cp.cuda.runtime.getDeviceCount()):
                if gpu_name.lower() in cp.cuda.runtime.getDeviceProperties(i)['name'].decode().lower():
                    d = cp.cuda.Device(i).use()
                    break
        else:
            d = cp.cuda.Device(0).use()
        if d is None:
            if gpu_name is not None:
                raise RuntimeError("Selected backend is pycuda, but no device found (name=%s)" % gpu_name)
            else:
                raise RuntimeError("Selected backend is pycuda, but no device found")
        gpu_ctx_dic["cupy"] = d

        # TODO: The following somehow helps initialising cupy, not sure why it's useful.
        #  (some context auto-init...). Otherwise a cuLaunchKernel error occurs with
        #  the first transform.
        cupy_a = cp.array(np.zeros((128, 128), dtype=np.float32))
        cupy_a.sum()
    else:
        raise RuntimeError("init_ctx: unknown backend ", backend)


def cleanup_cu_ctx():
    # Is that really clean ?
    if has_pycuda:
        if cu_drv.Context is not None:
            while cu_drv.Context.get_current() is not None:
                cu_drv.Context.pop()


atexit.register(cleanup_cu_ctx)


def l2(a, b):
    """L2 norm"""
    return np.sqrt((abs(a - b) ** 2).sum() / (abs(a) ** 2).sum())


def li(a, b):
    """Linf norm"""
    return abs(a - b).max() / abs(a).max()


def test_accuracy(backend, shape, ndim, axes, dtype, inplace, norm, use_lut,
                  r2c=False, dct=False, dst=False,
                  gpu_name=None, opencl_platform=None, stream=None, queue=None, return_array=False,
                  init_array=None, verbose=False, colour_output=False, ref_long_double=True, order='C'):
    """
    Measure the FT accuracy by comparing to the result from scipy (if available), or numpy.

    :param backend: either 'pyopencl', 'pycuda' or 'cupy'
    :param shape: the shape of the array to test. If this is an inplace r2c, the
        fast-axis length must be even, and two extra values will be appended along x,
        so the actual transform shape is the one supplied
    :param ndim: the number of FFT dimensions. Can be None if axes is given
    :param axes: the transform axes. Supersedes ndim
    :param dtype: either np.complex64 or np.complex128, or np.float32/np.float64 for r2c & dct
    :param inplace: if True, make an inplace transform. Note that for inplace r2c transforms,
        the size for the last (x, fastest) axis must be even.
    :param norm: either 0, 1 or "ortho"
    :param use_lut: if True,1, False or 0, will trigger useLUT=1 or 0 for VkFFT.
        If None, the default VkFFT behaviour is used.
    :param r2c: if True, test an r2c transform. If inplace, the last dimension
        (x, fastest axis) must be even
    :param dct: either 1, 2, 3 or 4 to test different dct. Only norm=1 is can be
        tested (native scipy normalisation).
    :param dst: either 1, 2, 3 or 4 to test different dst. Only norm=1 is can be
        tested (native scipy normalisation).
    :param gpu_name: the name of the gpu to use. If None, the first available
        for the backend will be used.
    :param opencl_platform: the name of the OpenCL platform to use. If None, the first available
        will be used.
    :param stream: the cuda stream to use, or None
    :param queue: the opencl queue to use (mandatory for the 'pyopencl' backend)
    :param return_array: if True, will return the generated random array so it can be
        re-used for different parameters
    :param init_array: the initial  (numpy) random array to use (should be filled
        with uniform random numbers between +/-0.5 for both real and imaginary
        fields), to save time. The correct type will be applied.
        If None, a random array is generated.
    :param verbose: if True, print a 1-line info for both fft and ifft results
    :param colour_output: if True, use some colour to tag the quality of the accuracy
    :param ref_long_double: if True and scipy is available, long double precision
        will be used for the reference transform. Otherwise, this is ignored.
    :param order: either 'C' (default C-contiguous) or 'F' to test a different
        stride. Note that for the latter, a 3D transform on a 4D array will not
        be supported as the last transform axis would be on the 4th dimension
        (once ordered by stride).
    :return: a dictionary with (l2_fft, li_fft, l2_ifft, li_ifft, tol, dt_array,
        dt_app, dt_fft, dt_ifft, src_unchanged_fft, src_unchanged_ifft, tol_test, str),
        with the L2 and Linf normalised norms comparing pyvkfft's result with either
        numpy, scipy, the reference
        tolerance, and the times spent in preparing the initial random array, creating
        the VkFFT app, and performing the forward and backward transforms (including
        the GPU and reference transforms, plus the L2 and Linf computations - don't use
        this for benchmarking), 'src_fft_unchanged' and 'srf_ifft_unchanged' are True
        if for an out-of-place transform, the source array is actually unmodified
        (which is not true for r2c ifft with ndim>=2).
        The last fields are 'tol_test' which is True if both li_fft and li_ifft are
        smaller than tol, and str the string summarising the results (printed if
        verbose is True). If return_array is True, the initial random array used is
        returned as 'd0'.
        All input parameters are also returned as key/values, except stream, queue,
        return_array, ini_array and verbose.
    """
    ndims = len(shape)
    if backend == "cupy" and has_cupy:
        mempool = cp.get_default_memory_pool()
        if mempool is not None:  # Is that test necessary ?
            # Clean memory pool, we are changing array sizes constantly, and using
            # N parallel process so memory management  must be done manually
            mempool.free_all_blocks()
    t0 = timeit.default_timer()
    init_ctx(backend, gpu_name=gpu_name, opencl_platform=opencl_platform, verbose=False)
    if backend == "pyopencl" and queue is None:
        queue = gpu_ctx_dic["pyopencl"][2]
    shape0 = shape
    dtype0 = dtype
    if dtype in (np.complex64, np.float32):
        dtype = np.complex64
        dtypef = np.float32
    else:
        dtype = np.complex128
        dtypef = np.float64

    if dct or dst:
        if norm != 1:
            raise RuntimeError("test_accuracy: only norm=1 can be used with dct or dst")

    # Need the fast axis for R2C (last for 'C' order, first for 'F')
    # Careful with special cases of axes with size 1, where stride does not increase or is 0...
    fast_axis = 0 if order == 'F' else ndims - 1

    if r2c:
        r2c_odd = shape[-1] % 2 == 1 if order == 'C' else shape[0] % 2 == 1
        r2c_inplace_pad = 1 if r2c_odd else 2
        if inplace:
            # Add two (or one for an odd-sized fast axis) extra columns in the source array
            # so the transform has the desired shape
            shape = list(shape)
            if order == 'C':
                shape[-1] += r2c_inplace_pad
            else:
                shape[0] += r2c_inplace_pad
        else:
            shapec = list(shape)
            shapec[fast_axis] = shapec[fast_axis] // 2 + 1
            shapec = tuple(shapec)
    else:
        r2c_odd = None
        r2c_inplace_pad = None
        shapec = tuple(shape)
    shape = tuple(shape)

    if init_array is not None:
        if r2c:
            if inplace:
                d0 = np.empty(shape, dtype=dtypef)
                if order == 'C':
                    d0[..., :-r2c_inplace_pad] = init_array
                else:
                    d0[:-r2c_inplace_pad] = init_array
            else:
                d0 = init_array.astype(dtypef)
        elif dct or dst:
            d0 = init_array.astype(dtypef)
        else:
            d0 = init_array.astype(dtype)
    else:
        rng = np.random.default_rng(seed=None)
        if r2c or dct or dst:
            d0 = rng.uniform(-0.5, 0.5, size=shape).astype(dtypef)
        else:
            d0 = (rng.uniform(-0.5, 0.5, size=shape) + 1j * rng.uniform(-0.5, 0.5, size=shape)).astype(dtype)

    if order != 'C':
        d0 = np.asarray(d0, order=order)

    t1 = timeit.default_timer()

    if 'opencl' in backend:
        app = clVkFFTApp(d0.shape, d0.dtype, queue, ndim=ndim, norm=norm,
                         axes=axes, useLUT=use_lut, inplace=inplace,
                         r2c=r2c, dct=dct, dst=dst, strides=d0.strides,
                         r2c_odd=r2c_odd)
        t2 = timeit.default_timer()
        d_gpu = cla.to_device(queue, d0)
        empty_like = cla.empty_like
    else:
        if backend == "pycuda":
            to_gpu = cua.to_gpu
            empty_like = cua.empty_like
        else:
            to_gpu = cp.array
            empty_like = cp.empty_like

        app = cuVkFFTApp(d0.shape, d0.dtype, ndim=ndim, norm=norm, axes=axes,
                         useLUT=use_lut, inplace=inplace, r2c=r2c,
                         dct=dct, dst=dst, strides=d0.strides,
                         r2c_odd=r2c_odd, stream=stream)
        t2 = timeit.default_timer()
        d_gpu = to_gpu(d0)

    if axes is None:
        axes_numpy = list(range(ndims))[-ndim:] if order == 'C' else list(range(ndim))
    else:
        # Make sure axes indices are >0
        axes_numpy = [ax if ax >= 0 else ndims + ax for ax in axes]

    if r2c:
        if fast_axis not in axes_numpy:
            raise RuntimeError(f"The fast axis must be transformed for R2C: "
                               f"fast axis={fast_axis}, axes_numpy={axes_numpy}")

        # For R2C, we need the same fast axis as on the GPU, or the
        # half-hermitian result won't look the same
        if fast_axis != axes_numpy[-1]:
            axes_numpy.remove(fast_axis)
            axes_numpy.append(fast_axis)
        # if order != 'C':
        #     print("R2C", ndims, ndim, shape, axes, axes_numpy, fast_axis, inplace)

    # base FFT scale for numpy (not used for DCT/DST)
    s = np.sqrt(np.prod([d0.shape[i] for i in axes_numpy]))
    if r2c and inplace:
        s = np.sqrt(s ** 2 / d0.shape[fast_axis] * (d0.shape[fast_axis] - r2c_inplace_pad))

    # Tolerance estimated from accuracy notebook
    if dtype in (np.complex64, np.float32):
        tol = 2e-6 + 5e-7 * np.log10(s ** 2)
    else:
        tol = 5e-15 + 5e-16 * np.log10(s ** 2)

    n = max(shape)
    bluestein = max(primes(n)) > 13
    if bluestein:
        tol *= 2
    ############################################################
    # FFT
    ############################################################
    if inplace:
        d1_gpu = d_gpu
    else:
        if r2c:
            if backend == "pyopencl":
                d1_gpu = cla.empty(queue, shapec, dtype=dtype, order=order)
            elif backend == "pycuda":
                d1_gpu = cua.empty(shapec, dtype=dtype, order=order)
            elif backend == "cupy":
                d1_gpu = cp.empty(shapec, dtype=dtype, order=order)
        else:
            # Note that pycuda's gpuarray.copy (as of 2022.2.2) does not copy strides
            d1_gpu = empty_like(d_gpu)

    if has_scipy and ref_long_double:
        # Use long double precision
        if r2c or dct or dst:
            d0n = d0.astype(np.longdouble)
        else:
            d0n = d0.astype(np.clongdouble)
    else:
        d0n = d0

    d1_gpu = app.fft(d_gpu, d1_gpu)
    if not (dct or dst):
        d1_gpu *= app.get_fft_scale()

    if r2c:
        if inplace:
            # Need to cut the fastest axis by 1 or 2
            d = rfftn(np.take(d0n, range(d0n.shape[fast_axis] - r2c_inplace_pad),
                              axis=fast_axis), axes=axes_numpy) / s
        else:
            d = rfftn(d0n, axes=axes_numpy) / s
    elif dct:
        d = dctn(d0n, axes=axes_numpy, type=dct)
    elif dst:
        d = dstn(d0n, axes=axes_numpy, type=dst)
    else:
        d = fftn(d0n, axes=axes_numpy) / s

    if inplace and r2c:
        assert d1_gpu.dtype == dtype, "The array type is incorrect after an inplace FFT"

    if r2c and d.shape != d1_gpu.shape:
        raise RuntimeError(f"Shape mismatch after the R2C FFT: shape={shape} shapec={shapec} "
                           f"order={order} axes={axes} axes_numpy={axes_numpy} ndim={ndim}")
    n2, ni = l2(d, d1_gpu.get()), li(d, d1_gpu.get())

    src_unchanged_fft = np.all(np.equal(d_gpu.get(), d0))

    # Output string
    if r2c:
        t = "R2C"
    elif dct:
        t = "DCT%d" % dct
    elif dst:
        t = "DST%d" % dst
    else:
        t = "C2C"
    # if r2c and inplace:
    #     tmp = list(d0.shape)
    #     if order == 'C':
    #         tmp[-1] -= r2c_inplace_pad
    #         shstr = str(tuple(tmp)).replace(" ", "")
    #         if ",)" in shstr:
    #             shstr = shstr.replace(",)", f"+{r2c_inplace_pad})")
    #         else:
    #             shstr = shstr.replace(")", f"+{r2c_inplace_pad})")
    #     else:
    #         tmp[0] -= r2c_inplace_pad
    #         if len(tmp) == 1:
    #             shstr = f"({tmp[0]}+{r2c_inplace_pad}),"
    #         else:
    #             shstr = f"({tmp[0]}+{r2c_inplace_pad},"
    #             shstr += str(tuple(tmp[1:])).replace(" ", "").replace(",)", ")")[1:]
    # else:
    #     shstr = str(d0.shape).replace(" ", "")
    shstr = app.get_shape_str()
    shax = str(axes).replace(" ", "").replace("[", "").replace("]", "")
    if colour_output:
        red = max(0, min(int((ni / tol - 0.2) * 255), 255))
        stol = "\x1b[48;2;%d;0;0m%5.1e < %5.1e (%5.3f)\x1b[0m" % (red, ni, tol, ni / tol)
    else:
        stol = "%5.1e < %5.1e (%5.3f)" % (ni, tol, ni / tol)
    nupstr = ''.join(str(nup) for nup in app.get_nb_upload())
    verb_out = "%8s %4s %15s axes=%12s ndim=%4s %5s %5s %10s lut=%4s inplace=%d " \
               " norm=%4s %s %5s: n2=%5.1e ninf=%s %d" % \
               (backend, t, shstr, shax, str(ndim), app.get_algo_str(), nupstr, str(d0.dtype),
                str(use_lut), int(inplace), str(norm), order, "FFT", n2, stol, src_unchanged_fft)

    t3 = timeit.default_timer()

    # Clean memory
    del d_gpu, d1_gpu
    if backend == "cupy" and has_cupy:
        mempool = cp.get_default_memory_pool()
        if mempool is not None:  # Is that test necessary ?
            # Clean memory pool, we are changing array sizes constantly, and using
            # N parallel process so memory management  must be done manually
            mempool.free_all_blocks()

    ############################################################
    # IFFT - from original array to avoid error propagation
    ############################################################
    if r2c:
        # Exception: we need a proper half-Hermitian array
        d0 = d.astype(dtype)
        if has_scipy and ref_long_double:
            d0n = d0.astype(np.clongdouble)
        else:
            d0n = d0
        if (np.argmin(strides_nonzero(d0.strides)[::-1]) != 0) and order == 'C':
            # np.fft.rfftn can change the fast axis
            d0 = np.asarray(d0, order='C')
    if order != 'C':
        d0 = np.asarray(d0, order=order)
    if 'opencl' in backend:
        d_gpu = cla.to_device(queue, d0)
    else:
        d_gpu = to_gpu(d0)

    if inplace:
        d1_gpu = d_gpu
    else:
        if r2c:
            if backend == "pyopencl":
                d1_gpu = cla.empty(queue, shape, dtype=dtypef, order=order)
            elif backend == "pycuda":
                d1_gpu = cua.empty(shape, dtype=dtypef, order=order)
            elif backend == "cupy":
                d1_gpu = cp.empty(shape, dtype=dtypef, order=order)
        else:
            d1_gpu = empty_like(d_gpu)

    d1_gpu = app.ifft(d_gpu, d1_gpu)
    if not (dct or dst):
        d1_gpu *= app.get_ifft_scale()

    if r2c:
        # We need to specify the destination shape, in the case
        # we want an odd-sized fast axis
        d = irfftn(d0n, s=[shape0[i] for i in axes_numpy], axes=axes_numpy) * s
    elif dct:
        d = idctn(d0n, axes=axes_numpy, type=dct)
    elif dst:
        d = idstn(d0n, axes=axes_numpy, type=dst)
    else:
        d = ifftn(d0n, axes=axes_numpy) * s

    if inplace:
        if dct or dst or r2c:
            assert d1_gpu.dtype == dtypef, "The array type is incorrect after an inplace iFFT"
        else:
            assert d1_gpu.dtype == dtype, "The array type is incorrect after an inplace iFFT"

    if r2c and inplace:
        tmp = np.take(d1_gpu.get(), range(d1_gpu.shape[fast_axis] - r2c_inplace_pad), axis=fast_axis)
        n2i, nii = l2(d, tmp), li(d, tmp)
    else:
        n2i, nii = l2(d, d1_gpu.get()), li(d, d1_gpu.get())

    src_unchanged_ifft = np.all(np.equal(d_gpu.get(), d0))

    # Max N for radix 1D C2R transforms to not overwrite source
    if platform.system() == 'Darwin':
        nmaxr2c1d = 2048 * (1 + int(dtype in (np.float32, np.complex64)))
    else:
        nmaxr2c1d = 3072 * (1 + int(dtype in (np.float32, np.complex64)))
    ndimtmp = ndim if ndim is not None else len(axes)
    if max(ni, nii) <= tol and (inplace or src_unchanged_fft) and \
            (inplace or src_unchanged_ifft or (r2c and ndimtmp > 1 or n >= nmaxr2c1d or bluestein)):
        success = 'OK'
    else:
        success = 'FAIL'

    if colour_output:
        red = max(0, min(int((nii / tol - 0.2) * 255), 255))
        stol = "\x1b[48;2;%d;0;0m%5.1e < %5.1e (%5.3f)\x1b[0m" % (red, nii, tol, nii / tol)
    else:
        stol = "%5.1e < %5.1e (%5.3f)" % (nii, tol, nii / tol)
    verb_out += "%5s: n2=%5.1e ninf=%s %d" % ("iFFT", n2i, stol, src_unchanged_ifft)

    # Also print the size of the allocated buffer and success
    verb_out += f" buf={app.get_tmp_buffer_str()} {success:4s}"

    if verbose:
        print(verb_out)

    t4 = timeit.default_timer()

    if backend == "pyopencl":
        gpu_name = gpu_ctx_dic["pyopencl"][0].name
    elif backend == "pycuda":
        gpu_name = gpu_ctx_dic["pycuda"][0].name()
    else:
        gpu_name = ""

    res = {"n2": n2, "ni": ni, "n2i": n2i, "nii": nii, "tol": tol, "dt_array": t1 - t0, "dt_app": t2 - t1,
           "dt_fft": t3 - t2, "dt_ifft": t4 - t3, "src_unchanged_fft": src_unchanged_fft,
           "src_unchanged_ifft": src_unchanged_ifft, "tol_test": max(ni, nii) < tol, "str": verb_out,
           "backend": backend, "shape": shape0, "ndim": ndim, "axes": axes, "dtype": dtype0, "inplace": inplace,
           "norm": norm, "use_lut": use_lut, "r2c": r2c, "dct": dct, "dst": dst,
           "gpu_name": gpu_name, "order": order}

    if return_array:
        res["d0"] = d0
    return res


def test_accuracy_kwargs(kwargs):
    # This function must be defined here, so it can be used with a multiprocessing pool
    # in test_fft, otherwise this will fail, see:
    # https://stackoverflow.com/questions/41385708/multiprocessing-example-giving-attributeerror
    if kwargs['backend'] == 'pyopencl' and has_opencl:
        try:
            t = test_accuracy(**kwargs)
        except cl.RuntimeError as ex:
            # The cl.RuntimeError can't be pickled, so is not correctly reported
            # when using multiprocessing. So we raise another and the traceback
            # from the previous one is still printed.
            raise RuntimeError("An OpenCL RuntimeError was encountered")
        return t
    return test_accuracy(**kwargs)


def exhaustive_test(backend, vn, ndim, dtype, inplace, norm, use_lut, r2c=False,
                    dct=False, dst=False, nproc=None,
                    verbose=True, return_res=False):
    """
    Run tests on a large range of sizes using multiprocessing. Manual function.

    :param backend: either 'pyopencl', 'pycuda' or 'cupy'
    :param vn: the list/iterable of sizes n.
    :param ndim: the number of dimensions. The array shape will be [n]*ndim
    :param dtype: either np.complex64 or np.complex128, or np.float32/np.float64 for r2c/dct/dst
    :param inplace: True or False
    :param norm: either 0, 1 or "ortho"
    :param use_lut: if True,1, False or 0, will trigger useLUT=1 or 0 for VkFFT.
        If None, the default VkFFT behaviour is used. Always True by default
        for double precision, so no need to force it.
    :param r2c: if True, test an r2c transform. If inplace, the last dimension
        (x, fastest axis) must be even
    :param dct: either 1, 2, 3 or 4 to test different dct. Only norm=1 is can be
        tested (native scipy normalisation).
    :param dst: either 1, 2, 3 or 4 to test different dst. Only norm=1 is can be
        tested (native scipy normalisation).
    :param nproc: the maximum number of parallel process to use. If None, the
        number of detected cores will be used (this may use too much memory !)
    :param verbose: if True, prints 1 line per test
    :param return_res: if True, return the list of result dictionaries.
    :return: True if all tests passed, False otherwise. If return_res is True, return
        the list of result dictionaries instead.
    """
    try:
        # Get the real number of processor cores available
        # os.sched_getaffinity is only available on some *nix platforms
        nproc1 = len(os.sched_getaffinity(0)) * psutil.cpu_count(logical=False) // psutil.cpu_count(logical=True)
    except AttributeError:
        nproc1 = os.cpu_count()
    if nproc is None:
        nproc = nproc1
    else:
        nproc = min(nproc, nproc1)
    # Generate the list of configurations as kwargs for test_accuracy()
    vkwargs = []
    for n in vn:
        kwargs = {"backend": backend, "shape": [n] * ndim, "ndim": ndim, "axes": None, "dtype": dtype,
                  "inplace": inplace, "norm": norm, "use_lut": use_lut,
                  "r2c": r2c, "dct": dct, "dst": dst, "stream": None,
                  "verbose": False}
        vkwargs.append(kwargs)
    vok = []
    vres = []
    # Need to use spawn to handle the GPU context
    with multiprocessing.get_context('spawn').Pool(nproc) as pool:
        for res in pool.imap(test_accuracy_kwargs, vkwargs):
            # TODO: this should better be logged
            if verbose:
                print(res['str'])
            ni, n2 = res["ni"], res["n2"]
            nii, n2i = res["nii"], res["n2i"]
            tol = res["tol"]
            ok = max(ni, nii) < tol
            if not inplace:
                ok = ok and res["src_unchanged_fft"]
                if not r2c:
                    ok = ok and res["src_unchanged_ifft"]
            vok.append(ok)
            if return_res:
                vres.append(res)
    if return_res:
        return vres
    return np.all(vok)
