#pylint: disable=no-name-in-module

import numpy as np
from mpi4py import MPI
from .utilities import FFTW_FORWARD, FFTW_MEASURE

def get_fftw_lib(dtype):
    """Return compiled fftw module interfacing the FFTW library

    Parameters
    ----------
    dtype : dtype
        Data precision

    Returns
    -------
    Module or ``None``
        Module can be either :mod:`.fftwf_xfftn`, :mod:`.fftw_xfftn` or
        :mod:`.fftwl_xfftn`, depending on precision.
    """

    dtype = np.dtype(dtype).char.upper()
    if dtype == 'G':
        try:
            from . import fftwl_xfftn
            return fftwl_xfftn
        except ImportError: #pragma: no cover
            return None
    elif dtype == 'D':
        try:
            from . import fftw_xfftn
            return fftw_xfftn
        except ImportError: #pragma: no cover
            return None
    elif dtype == 'F':
        try:
            from . import fftwf_xfftn
            return fftwf_xfftn
        except ImportError: #pragma: no cover
            return None
    else: #pragma: no cover
        return None

fftlib = {}
for t in 'fdg':
    fftw_lib = get_fftw_lib(t)
    if fftw_lib is not None:
        fftlib[t.upper()] = fftw_lib

comm = MPI.COMM_WORLD

def get_planned_FFT(input_array, output_array, axes=(-1,), kind=FFTW_FORWARD,
                    threads=1, flags=(FFTW_MEASURE,), normalization=1.0):
    """Return instance of transform class

    Parameters
    ----------
    input_array : array
        real or complex input array
    output_array : array
        real or complex output array
    axes : sequence of ints, optional
        The axes to transform over, starting from the last
    kind : int or sequence of ints, optional
        Any one of (or possibly several for real-to-real)

            - FFTW_FORWARD (-1)
            - FFTW_R2HC (0)
            - FFTW_BACKWARD (1)
            - FFTW_HC2R (1)
            - FFTW_DHT (2)
            - FFTW_REDFT00 (3)
            - FFTW_REDFT01 (4)
            - FFTW_REDFT10 (5)
            - FFTW_REDFT11 (6)
            - FFTW_RODFT00 (7)
            - FFTW_RODFT01 (8)
            - FFTW_RODFT10 (9)
            - FFTW_RODFT11 (10)
    threads : int, optional
        Number of threads to use in transforms
    flags : int or sequence of ints, optional
        Any one of, but not necessarily for all transforms or all combinations

            - FFTW_MEASURE (0)
            - FFTW_DESTROY_INPUT (1)
            - FFTW_UNALIGNED (2)
            - FFTW_CONSERVE_MEMORY (4)
            - FFTW_EXHAUSTIVE (8)
            - FFTW_PRESERVE_INPUT (16)
            - FFTW_PATIENT (32)
            - FFTW_ESTIMATE (64)
            - FFTW_WISDOM_ONLY (2097152)
    normalization : float, optional
        Normalization factor

    Returns
    -------
    :class:`.fftwf_xfftn.FFT`, :class:`.fftw_xfftn.FFT` or :class:`.fftwl_xfftn.FFT`
        An instance of the return type configured for the desired transforms

    """
    dtype = input_array.dtype.char
    assert dtype.upper() in fftlib
    _fft = fftlib[dtype.upper()]
    return _fft.FFT(input_array, output_array, axes, kind, threads, flags,
                    normalization)

def export_wisdom(filename):
    """Export FFTW wisdom

    Parameters
    ----------
    filename : str
        Name of file used to export wisdom to

    Note
    ----
    Wisdom is stored for all precisions available: float, double and long
    double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
    n is the rank of the processor.
    Wisdom is imported using :func:`.import_wisdom`, which must be called
    with the same MPI configuration as used with :func:`.export_wisdom`.

    See also
    --------
    :func:`.import_wisdom`

    """
    rank = str(comm.Get_rank())
    e = []
    for key, lib in fftlib.items():
        e.append(lib.export_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
    assert np.all(np.array(e) == 1), "Not able to export wisdom {}".format(filename)

def import_wisdom(filename):
    """Import FFTW wisdom

    Parameters
    ----------
    filename : str
        Name of file used to import wisdom from

    Note
    ----
    Wisdom is imported for all available precisions: float, double and long
    double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
    n is the rank of the processor.
    Wisdom is exported using :func:`.export_wisdom`.
    Note that importing wisdom only works when using the same MPI configuration
    as used with :func:`.export_wisdom`.


    See also
    --------
    :func:`.export_wisdom`

    """
    rank = str(comm.Get_rank())
    e = []
    for key, lib in fftlib.items():
        e.append(lib.import_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
    assert np.all(np.array(e) == 1), "Not able to import wisdom {}".format(filename)

def forget_wisdom():
    for lib in fftlib.values():
        lib.forget_wisdom()

def set_timelimit(limit):
    """Set time limit for planning

    Parameters
    ----------
    limit : number
        The new time limit set for planning of serial transforms
    """
    for lib in fftlib.values():
        lib.set_timelimit(limit) # limit's precision handled by cython

def cleanup():
    for lib in fftlib.values():
        lib.cleanup()