# accel_math.py
#
# Various functions related to accelerated computations using FFTW, CUDA, numexpr, and related.
#
import numpy as np
import multiprocessing
import matplotlib.pyplot as plt
from . import conf

import time
import logging
_log = logging.getLogger('poppy')


try:
    # try to import FFTW to see if it is available
    import pyfftw
    # Setup infrastructure for FFTW
    _FFTW_INIT = {}  # dict of array sizes for which we have already performed the required FFTW planning step
    _FFTW_FLAGS = ['measure']
    _FFTW_AVAILABLE = True
except ImportError:
    pyfftw = None
    _FFTW_AVAILABLE = False

try:
    # try to import numexpr package to see if it is available
    import numexpr as ne

    _NUMEXPR_AVAILABLE = True
except ImportError:
    ne = None
    _NUMEXPR_AVAILABLE = False


try:
    import mkl_fft
    _MKLFFT_AVAILABLE = True
except ImportError:
    _MKLFFT_AVAILABLE = False


try:
    # try to import CUDA packages to see if they are available
    import pyculib
    from numba import cuda
    _CUDA_PLANS = {} # plans for various array sizes already prepared
    _CUDA_AVAILABLE = True
except ImportError:
    pyculib = None
    _CUDA_AVAILABLE = False

try:
    # try to import OpenCL packages to see if they are is available
    import pyopencl
    import pyopencl.array
    import gpyfft
    _OPENCL_AVAILABLE = True
    _OPENCL_STATE = dict()
except ImportError:
    _OPENCL_AVAILABLE = False

_USE_CUDA = (conf.use_cuda and _CUDA_AVAILABLE)
_USE_OPENCL = (conf.use_opencl and _OPENCL_AVAILABLE)
_USE_NUMEXPR = (conf.use_numexpr and _NUMEXPR_AVAILABLE)
_USE_FFTW = (conf.use_fftw and _FFTW_AVAILABLE)
_USE_MKL = (conf.use_mkl and _MKLFFT_AVAILABLE)


def update_math_settings():
    """ Update the module-level math flags, based on user settings
    """
    global _USE_CUDA, _USE_OPENCL, _USE_NUMEXPR, _USE_FFTW, _USE_MKL
    _USE_CUDA = (conf.use_cuda and _CUDA_AVAILABLE)
    _USE_OPENCL = (conf.use_opencl and _OPENCL_AVAILABLE)
    _USE_NUMEXPR = (conf.use_numexpr and _NUMEXPR_AVAILABLE)
    _USE_FFTW = (conf.use_fftw and _FFTW_AVAILABLE)
    _USE_MKL = (conf.use_mkl and _MKLFFT_AVAILABLE)


def _float():
    """ Returns numpy data type for desired precision based on configuration """
    # How many bits per float to use?
    return np.float64 if conf.double_precision else np.float32


def _complex():
    """ Returns numpy data type for desired precision based on configuration """
    # How many bits per complex float to use?
    return np.complex128 if conf.double_precision else np.complex64


def _r(x, y):
    """ Function to speed up computing the radius given x and y, using Numexpr if available
    Otherwise defaults to numpy. """
    if _USE_NUMEXPR:
        return ne.evaluate("sqrt(x**2+y**2)")
    else:
        return np.sqrt(x ** 2 + y ** 2)


def _exp(x):
    """
    Function to speed up taking exponential of an array if NumExpr is available.
    Otherwise defaults to np.exp()

    """
    if _USE_NUMEXPR:
        return ne.evaluate("exp(x)", optimization='moderate', )
    else:
        return np.exp(x)

def _fftshift(x):
    """ FFT shifts of array contents, using CUDA if available.
    Otherwise defaults to numpy.

    Note - TODO write an OpenCL version

    See also ifftshift
    """

    N=x.shape[0]
    # the CUDA fftshift is set up to work on blocks of 32, so
    # N must be a multiple of 32. We check this rapidly using a bit mask:
    #    (x & 31)==0  is a ~20x faster equivalent of (np.mod(x,32)==0)
    if (_USE_CUDA) & (N==x.shape[1]) & ((N & 31)==0):
        blockdim = (32, 32) # threads per block
        numBlocks = (int(N/blockdim[0]),int(N/blockdim[1]))
        cufftShift_2D_kernel[numBlocks, blockdim](x.ravel(),N)
        return x
    else:
        return np.fft.fftshift(x)

def _ifftshift(x):
    """ Inverse FFT shifts of array contents, using CUDA if available.
    Otherwise defaults to numpy.
    Note, ifftshift and fftshift are identical for even-length x,
    the functions differ by one sample for odd-length x. This function
    implicitly assumes that if using CUDA, the array size must be even,
    so we can use the same algorithm as fftshift.

    Note - TODO write an OpenCL version

    See also fftshift
    """

    N=x.shape[0]
    # the CUDA fftshift is set up to work on blocks of 32, so
    # N must be a multiple of 32. We check this rapidly using a bit mask:
    #   not (x & 31)  is a ~20x faster equivalent of (np.mod(x,32)==0)
    if (_USE_CUDA) & (N==x.shape[1]) & ((N & 31)==0):
        blockdim = (32, 32) # threads per block
        numBlocks = (int(N/blockdim[0]),int(N/blockdim[1]))
        cufftShift_2D_kernel[numBlocks, blockdim](x.ravel(),N)
        return x
    else:
        return np.fft.ifftshift(x)




def fft_2d(wavefront, forward=True, normalization=None, fftshift=True):
    """ main entry point for FFTs, used in Wavefront._propagate_fft and
    elsewhere. This will invoke one of the following, depending on availability:
        - CUDA on NVidia GPU
        - OpenCL on AMD GPU
        - FFTW on CPU
        - numpy on CPU

    This function handles ONLY the core numerics itself, as fast as possible,
    (and some minor related logging) .
    All the interaction with object state for Wavefront arrays should happen elsewhere.

    TODO: this should execute an IN PLACE FFT, so we don't have to pass around arrays to return
    anything.

    Parameters
    -----------
    forward : bool
        set to True for forward FFT, False for inverse fft
        SIGN CONVENTION: Note, forward optical propagation requires an inverse FFT, and
        vice versa. This is to have a positive sign in the exponent for forward propagations.
        See sign convention documentation.
    normalization : float, optional
        Normalization factor. Defaults to 1./wavefront.shape[0] for forward,
        and wavefront.shape[0] for inverse. Use this only if you need a non-default
        behavior.
    fftshift : bool
        apply FFT shift after forwards propagation or before inverse propagation?
        SIGN CONVENTION: Note, that means after an inverse FFT or before a forward
        FFT. This differs from the behavior of poppy versions < 1.0.

    """
    ## To use a fast FFT, it must both be enabled and the library itself has to be present
    global _USE_OPENCL, _USE_CUDA # need to declare global in case we need to change it, below
    t0 = time.time()

    # OpenCL cfFFT only can FFT certain array sizes.
    if _USE_OPENCL and not isproductofsmallprimes(wavefront.shape[0]):
        _log.debug(("Wavefront size {} not supported by OpenCL, therefore disabling "+
            "USE_OPENCL for this calculation.").format(wavefront.shape))
        _USE_OPENCL = False

    # This annoyingly complicated if/elif is just for the debug print statement
    if _USE_CUDA:
        method = 'pyculib (CUDA GPU)'
    elif _USE_OPENCL:
        method = 'pyopencl (OpenCL GPU)'
    elif _USE_MKL:
        method = 'mkl_fft'
    elif _USE_FFTW:
        method = 'pyfftw'
    else:
        method = 'numpy'
    _log.debug("using {2} FFT of {0} array, FFT_direction={1}".format(
        str(wavefront.shape), 'forward' if forward else 'backward', method))

    if (forward) and fftshift: # shift before backwards propagations (using forward FFT)
        # This needs to be an ifftshift, for machine-precision equivalence to MFT
        # as tested by test_MFT_FFT_equivalence_in_OpticalSystem; ifftshift and fftshift are
        # precisely equivalent for even-sized arrays but differ by 1 pixel offset for odd-sized
        # arrays. The use of ifftshift here ensures the PSF is correctly centered in the center pixel
        # for odd-sized arrays.
        wavefront = _ifftshift(wavefront)

    t1 = time.time()
    if _USE_CUDA:
        if normalization is None:
            normalization = 1./wavefront.shape[0]  # regardless of direction, for CUDA

        # We need a CUDA FFT plan for each size and shape of FFT.
        # The plans can be cached for reuse, since they cost some
        # 10s of milliseconds to create
        params = (wavefront.shape, wavefront.dtype, wavefront.dtype)
        try:
            cufftplan = _CUDA_PLANS[params]
        except KeyError:
            cufftplan = pyculib.fft.FFTPlan(*params)
            _CUDA_PLANS[params] = cufftplan

        # perform FFT on GPU, and return results in place to same array.
        if forward:
            cufftplan.forward(wavefront, out=wavefront)
        else:
            cufftplan.inverse(wavefront, out=wavefront)

    elif _USE_OPENCL:
        if normalization is None:
            normalization = 1./wavefront.shape[0] if forward else wavefront.shape[0]

        context, queue = get_opencl_context()
        wf_on_gpu = pyopencl.array.to_device(queue, wavefront)
        transform = gpyfft.fft.FFT(context, queue, wf_on_gpu, axes=(0,1))
        event, = transform.enqueue(forward=forward)
        event.wait()
        wavefront[:] = wf_on_gpu.get()
        del wf_on_gpu

    elif _USE_MKL:
        # Intel MKL is a drop-in replacement for numpy fft but much faster
        do_fft = mkl_fft.fft2 if forward else mkl_fft.ifft2
        if normalization is None:
            normalization = 1./wavefront.shape[0] if forward else wavefront.shape[0]
        wavefront = do_fft(wavefront)

    elif _USE_FFTW:
        FFT_direction = 'forward' if forward else 'backward' # back compatible for use in _FFTW_INIT
        do_fft = pyfftw.interfaces.numpy_fft.fft2 if forward else pyfftw.interfaces.numpy_fft.ifft2
        if normalization is None:
            normalization = 1./wavefront.shape[0] if forward else wavefront.shape[0]


        if (wavefront.shape, FFT_direction) not in _FFTW_INIT:
            # The first time you run FFTW to transform a given size, it does a speed test to
            # determine optimal algorithm that is destructive to your chosen array.
            # So only do that test on a copy, not the real array:
            _log.info("Measuring pyfftw optimal plan for %s, direction=%s" % (
                str(wavefront.shape), FFT_direction))

            pyfftw.interfaces.cache.enable()
            pyfftw.interfaces.cache.set_keepalive_time(30)

            test_array = np.zeros(wavefront.shape)
            test_array = do_fft(test_array, overwrite_input=True, planner_effort='FFTW_MEASURE',
                                threads=multiprocessing.cpu_count())

            _FFTW_INIT[(wavefront.shape, FFT_direction)] = True

        wavefront = do_fft(wavefront, overwrite_input=True, planner_effort='FFTW_MEASURE',
                                threads=multiprocessing.cpu_count())
    else: # Basic numpy FFT
        do_fft =  np.fft.fft2 if forward else np.fft.ifft2
        if normalization is None:
            normalization = 1./wavefront.shape[0] if forward else wavefront.shape[0]
        wavefront = do_fft(wavefront)
    t2 = time.time()

    if (not forward) and fftshift:  # shift after forwards propagations (using inverse FFT)
        # This needs to be a fftshift here, since we use ifftshift above
        # See comment above in this function.
        wavefront = _fftshift(wavefront)

    wavefront *= normalization
    t3 = time.time()
    _log.debug("    FFT_2D: FFT in {:3f} s, full function  in {:.3f} s".format(t2-t1, t3-t0))

    return wavefront



def ispowerof2(num):
    """ Is this number a power of 2?"""
    # see http://code.activestate.com/recipes/577514-chek-if-a-number-is-a-power-of-two/
    return (num & (num-1) == 0)


def isproductofsmallprimes(num):
    """ CLFFT only supports array sizes which are products of primes <= 13;
    Is this integer a product of primes no greater than that?
    """
    tmp = num*1 # make copy
    small_primes = [2,3,5,7,9,11,13]
    for p in small_primes:
        quotient, remainder = divmod(tmp, p)
        while not remainder:
            tmp = quotient
            quotient, remainder = divmod(tmp, p)
        if tmp==1: return True # if we get to 1 we're done
    return False # if we don't get to 1 then it's got higher factors

if _OPENCL_AVAILABLE:
    def get_opencl_context():
        """ Create, save, and retrieve OpenCL handles to the GPU """
        if len(_OPENCL_STATE) == 0:
            platforms = pyopencl.get_platforms()
            if len(platforms) == 1:
                _OPENCL_STATE['platform'] = platforms[0]
            else:
                raise RuntimeError("OpenCL code needs update for multiple platforms")
            gpus = _OPENCL_STATE['platform'].get_devices(device_type=pyopencl.device_type.GPU)
            if len(gpus) == 1:
                device = gpus[0]
                _OPENCL_STATE['device'] = device
            else:
                #raise RuntimeError("OpenCL code could not uniquely identify which device to use as GPU")
                device = gpus[1]
                _log.warning('Caution - hard coded use of gpu #1 if > 1 GPUs present')
            context = pyopencl.Context(devices=[device])
            queue = pyopencl.CommandQueue(context)

            _OPENCL_STATE['context'] = context
            _OPENCL_STATE['queue'] = queue
        return (_OPENCL_STATE['context'], _OPENCL_STATE['queue'])





if  _CUDA_AVAILABLE:
    @cuda.jit()
    def cufftShift_2D_kernel(data, N):
        """
        adopted CUDA FFT shift code from:
        https://github.com/marwan-abdellah/cufftShift
        (GNU Lesser Public License)
        """

        # // 2D Slice & 1D Line
        sLine = N
        sSlice = N * N
        # // Transformations Equations
        sEq1 = int((sSlice + sLine) / 2)
        sEq2 = int((sSlice - sLine) / 2)
        x, y = cuda.grid(2)
        # // Thread Index Converted into 1D Index
        index = (y * N) + x

        if x < N / 2:
            if y < N / 2:
                # // First Quad
                temp = data[index]
                data[index] = data[index + sEq1]
                # // Third Quad
                data[index + sEq1] = temp
        else:
            if y < N / 2:
                # // Second Quad
                temp = data[index]
                data[index] = data[index + sEq2]
                data[index + sEq2] = temp

# ##################################################################
#
#     Performance benchmarks


def benchmark_fft(npix=2048, iterations=20, double_precision=True):
    """ Performance benchmark function for standard imaging """
    #import poppy
    import timeit
    import poppy


    complextype = 'complex128' if double_precision else 'complex64'

    timer = timeit.Timer("tmp2 = poppy.accel_math.fft_2d(tmp, fftshift=False)",
            setup="""
import poppy
import numpy as np
tmp = np.asarray(np.random.rand({npix},{npix}), np.{complextype})
            """.format(npix=npix, complextype=complextype))
    print("Timing performance of FFT for {npix} x {npix}, {complextype}, with {iterations} iterations".format(
        npix=npix, iterations=iterations, complextype=complextype))

    defaults = (poppy.conf.use_mkl, poppy.conf.use_fftw, poppy.conf.use_numexpr, poppy.conf.use_cuda,
            poppy.conf.use_opencl, poppy.conf.double_precision)
    poppy.conf.double_precision = double_precision

    # Time baseline performance in numpy
    print("Timing performance in plain numpy:")

    poppy.conf.use_mkl, poppy.conf.use_fftw, poppy.conf.use_numexpr, poppy.conf.use_cuda, poppy.conf.use_opencl = (False, False, False, False, False)
    update_math_settings()
    time_numpy = timer.timeit(number=iterations) / iterations
    print("  {:.3f} s".format(time_numpy))

    if poppy.accel_math._FFTW_AVAILABLE:
        print("Timing performance with FFTW:")
        poppy.conf.use_fftw = True
        update_math_settings()
        time_fftw = timer.timeit(number=iterations) / iterations
        print("  {:.3f} s".format(time_fftw))
    else:
        time_fftw = np.NaN

    if poppy.accel_math._NUMEXPR_AVAILABLE:
        print("Timing performance with Numexpr + FFTW:")
        poppy.conf.use_numexpr = True
        update_math_settings()
        time_numexpr = timer.timeit(number=iterations) / iterations
        print("  {:.3f} s".format(time_numexpr))
    else:
        time_numexpr = np.NaN

    if poppy.accel_math._MKLFFT_AVAILABLE:
        print("Timing performance with Numexpr + MKL:")
        poppy.conf.use_numexpr = True
        poppy.conf.use_fftw = False
        poppy.conf.use_mkl = True
        update_math_settings()
        time_mkl = timer.timeit(number=iterations) / iterations
        print("  {:.3f} s".format(time_mkl))
    else:
        time_mkl = np.NaN

    if poppy.accel_math._CUDA_AVAILABLE:
        print("Timing performance with CUDA:")
        poppy.conf.use_cuda = True
        poppy.conf.use_opencl = False
        update_math_settings()
        time_cuda = timer.timeit(number=iterations) / iterations
        print("  {:.3f} s".format(time_cuda))
    else:
        time_cuda = np.NaN

    if poppy.accel_math._OPENCL_AVAILABLE:
        print("Timing performance with OpenCL:")
        poppy.conf.use_opencl = True
        poppy.conf.use_cuda = False
        update_math_settings()
        time_opencl = timer.timeit(number=iterations) / iterations
        print("  {:.3f} s".format(time_opencl))
    else:
        time_opencl = np.NaN


    poppy.conf.use_mkl, poppy.conf.use_fftw, poppy.conf.use_numexpr, poppy.conf.use_cuda,\
            poppy.conf.use_opencl, poppy.conf.double_precision = defaults

    return {'numpy': time_numpy,
            'fftw': time_fftw,
            'numexpr': time_numexpr,
            'mkl': time_mkl,
            'cuda': time_cuda,
            'opencl': time_opencl}


def get_processor_name():
    """Return processor hardware name, in a cross-platform manner

    With thanks to https://stackoverflow.com/a/20161999
    """
    import os, platform, subprocess, re

    if platform.system() == "Windows":
        return platform.processor()
    elif platform.system() == "Darwin":
        os.environ['PATH'] = os.environ['PATH'] + os.pathsep + '/usr/sbin'
        command = "sysctl -n machdep.cpu.brand_string".split()
        return (subprocess.check_output(command).strip()).decode()
    elif platform.system() == "Linux":
        command = "cat /proc/cpuinfo"
        all_info = subprocess.check_output(command, shell=True).strip()
        for line in all_info.split("\n"):
            if "model name" in line:
                return re.sub(".*model name.*:", "", line, 1)
    return ""


def get_physical_cpu_count():
    """Get physical number of CPU cores

    You might think multiprocessing.cpu_count() would give this, but no.
    Hyperthreading makes that number harder to interpret.

    """
    import os, platform, subprocess, re

    if platform.system() == "Darwin":
        os.environ['PATH'] = os.environ['PATH'] + os.pathsep + '/usr/sbin'
        command = "sysctl -n hw.physicalcpu".split()
        return subprocess.check_output(command).strip().decode()
    else:
        raise NotImplementedError("make this cross platform")

def benchmark_2d_ffts(mode='poppy', max_pow=13, verbose=False, savefig=False):
    """Benchmark FFT runtime vs array size, for multiple transform libraries

    Parameters
    ----------
    mode : string, 'poppy' or 'base'
        What to test, either the full poppy usage of the 2D transform, including
        surrounding setup and normalization code, or the basic transform on its own.
    max_pow : int
        Maximum power of 2 array size to test up to
    verbose : bool
        output numpy config info
    savefig : bool
        save plot result to a PDF?

    """

    # modified based on https://github.com/numpy/numpy/issues/17839#issuecomment-733543850


    try:
        from simple_benchmark import benchmark
    except ImportError:
        raise RuntimeError("You must have simple_benchmark installed to run this function")

    threads = multiprocessing.cpu_count()

    if mode == 'base':
        # Functions for benchmarking the low-level FFT functions in each library
        def pocketfft(it):
            np.fft.fft2(it)

        def scipyfft(it):
            scipy.fft.fft2(it)

        def fftw_1thread(it):
            # note, FFTW defaults to 1 thread, unless you override the config,
            # but that's not a fair comparison
            pyfftw.interfaces.numpy_fft.fft2(it, thread=1)

        def fftw(it):
            # explicitly try multithreaded here
            pyfftw.interfaces.numpy_fft.fft2(it, threads=threads)

        def mklfft(it):
            mkl_fft.fft2(it)

        funcs_to_test = [pocketfft, fftw, mklfft, scipyfft, fftw_1thread, ]
        function_aliases = {pocketfft: 'numpy.fft', fftw: 'pyfftw.fft, multithreaded', scipyfft: 'scipy.fft',
                            mklfft: "MKL FFT", fftw_1thread: 'pyfftw.fft, single-thread'}
        title = "Basic 2D FFT only"

    elif mode == 'poppy':
        # Functions for benchmarking poppy transforms using each library

        def poppy_numpyfft(it):
            global _USE_FFTW, _USE_MKL
            _USE_FFTW = False
            _USE_MKL = False

            fft_2d(it, fftshift=False)

        def poppy_fftw(it):
            global _USE_FFTW, _USE_MKL
            _USE_FFTW = True
            _USE_MKL = False

            fft_2d(it, fftshift=False)

        def poppy_mklfft(it):
            global _USE_FFTW, _USE_MKL
            _USE_FFTW = False
            _USE_MKL = True

            fft_2d(it, fftshift=False)

        funcs_to_test = [poppy_numpyfft, poppy_fftw, poppy_mklfft]
        function_aliases = {poppy_numpyfft: 'poppy using numpy.fft', poppy_fftw: 'poppy using pyfftw.fft',
                            poppy_mklfft: "poppy using MKL FFT"}
        title = 'full poppy.accel_math.fft_2d'
    else:
        raise ValueError(f"Unknown/invalid value for 'base' parameter: {base}")

    def shp(len):
        return (len, len)

    if verbose:
        print(np.version)
        np.show_config()

    # Turn on the cache for optimum performance
    pyfftw.interfaces.cache.enable()
    b_array = benchmark(
        funcs_to_test,
        arguments={2 ** i: np.random.uniform(size=shp(2 ** i)) + 1j * np.random.uniform(size=shp(2 ** i)) for i in
                   range(2, max_pow)},
        argument_name='array size',
        function_aliases=function_aliases
    )
    plt.figure(figsize=(12, 8))
    b_array.plot()
    plt.grid(which='both', alpha=0.2)
    plt.grid(which='major', alpha=0.5)
    plt.xlim(1, 2e4)
    plt.ylim(1e-6, 1e1)

    cpu_label = get_processor_name() + f", {get_physical_cpu_count()} cores"
    plt.title(f"Time for 2D complex FFTs: {title}\n{cpu_label}", fontweight='bold')
    if savefig:
        plt.savefig(f"bench_ffts_{mode}.png")


def benchmark_2d_mfts(max_pow=13, savefig=False):
    """Benchmark MFT runtime vs array size, for multiple transform libraries

    Parameters
    ----------
    mode : string, 'poppy' or 'base'
        What to test, either the full poppy usage of the 2D transform, including
        surrounding setup and normalization code, or the basic transform on its own.
    max_pow : int
        Maximum power of 2 array size to test up to
    savefig : bool
        save plot result to a PDF?

    """

    # modified based on https://github.com/numpy/numpy/issues/17839#issuecomment-733543850
    from simple_benchmark import benchmark

    import functools
    import poppy

    def shp(len):
        return (len, len)

    def test_mft_numpy(array, npix=64):
        global _USE_NUMEXPR

        _USE_NUMEXPR = False
        poppy.matrixDFT.matrix_dft(array, 16, npix)

    def test_mft_numexpr(array, npix=64):
        global _USE_NUMEXPR
        _USE_NUMEXPR = True

        poppy.matrixDFT.matrix_dft(array, 16, npix)

    test_mft_numpy_512 = functools.partial(test_mft_numpy, npix=512)
    test_mft_numexpr_512 = functools.partial(test_mft_numexpr, npix=512)

    b_array = benchmark(
        [test_mft_numpy, test_mft_numexpr, test_mft_numpy_512, test_mft_numexpr_512],
        arguments={2 ** i: np.random.uniform(size=shp(2 ** i)) + 1j * np.random.uniform(size=shp(2 ** i)) for i in
                   range(2, max_pow)},
        argument_name='pupil array size, npupil',
        function_aliases={test_mft_numpy: "MFT with numpy, npix=64", test_mft_numexpr: "MFT with numexpr, npix=64",
                          test_mft_numpy_512: "MFT with numpy, npix=512",
                          test_mft_numexpr_512: "MFT with numexpr, npix=512"}
    )
    plt.figure(figsize=(12, 8))
    b_array.plot()
    plt.grid(which='both', alpha=0.2)
    plt.grid(which='major', alpha=0.5)
    plt.xlim(1, 2e4)
    plt.ylim(1e-6, 1e1)

    cpu_label = get_processor_name() + f", {get_physical_cpu_count()} cores"
    plt.title(f"Matrix Fourier Transform timings\n{cpu_label}", fontweight='bold')

    if savefig:
        plt.savefig(f"bench_mfts.png")