# Source code for mpi4py_fft.libfft

    import functools
    import numpy as np
    from . import fftw

    def _Xfftn_plan_pyfftw(shape, axes, dtype, transforms, options):

        import pyfftw
        opts = dict(
            avoid_copy=True,
            overwrite_input=True,
            auto_align_input=True,
            auto_contiguous=True,
            threads=1,
        )
        opts.update(options)

        transforms = {} if transforms is None else transforms
        if tuple(axes) in transforms:
            plan_fwd, plan_bck = transforms[tuple(axes)]
        else:
            if np.issubdtype(dtype, np.floating):
                plan_fwd = pyfftw.builders.rfftn
                plan_bck = pyfftw.builders.irfftn
            else:
                plan_fwd = pyfftw.builders.fftn
                plan_bck = pyfftw.builders.ifftn

        s = tuple(np.take(shape, axes))
        U = pyfftw.empty_aligned(shape, dtype=dtype)
        xfftn_fwd = plan_fwd(U, s=s, axes=axes, **opts)
        U.fill(0)
        if np.issubdtype(dtype, np.floating):
            del opts['overwrite_input']
        V = xfftn_fwd.output_array
        xfftn_bck = plan_bck(V, s=s, axes=axes, **opts)
        V.fill(0)
        xfftn_fwd.update_arrays(U, V)
        xfftn_bck.update_arrays(V, U)

        wrapped_xfftn_bck = functools.partial(xfftn_bck, normalise_idft=False)
        functools.update_wrapper(wrapped_xfftn_bck, xfftn_bck,
                                 assigned=['input_array',
                                           'output_array',
                                           '__doc__'])

        return (xfftn_fwd, wrapped_xfftn_bck)

    def _Xfftn_plan_fftw(shape, axes, dtype, transforms, options):

        opts = dict(
            overwrite_input='FFTW_DESTROY_INPUT',
            planner_effort='FFTW_MEASURE',
            threads=1,
        )
        opts.update(options)
        flags = (fftw.flag_dict[opts['planner_effort']],
                 fftw.flag_dict[opts['overwrite_input']])
        threads = opts['threads']

        transforms = {} if transforms is None else transforms
        if tuple(axes) in transforms:
            plan_fwd, plan_bck = transforms[tuple(axes)]
        else:
            if np.issubdtype(dtype, np.floating):
                plan_fwd = fftw.rfftn
                plan_bck = fftw.irfftn
            else:
                plan_fwd = fftw.fftn
                plan_bck = fftw.ifftn

        s = tuple(np.take(shape, axes))
        U = fftw.aligned(shape, dtype=dtype)
        xfftn_fwd = plan_fwd(U, s=s, axes=axes, threads=threads, flags=flags)
        U.fill(0)
        V = xfftn_fwd.output_array
        if np.issubdtype(dtype, np.floating):
            flags = (fftw.flag_dict[opts['planner_effort']],)
        xfftn_bck = plan_bck(V, s=s, axes=axes, threads=threads, flags=flags, output_array=U)
        return (xfftn_fwd, xfftn_bck)

    def _Xfftn_plan_numpy(shape, axes, dtype, transforms, options):

        transforms = {} if transforms is None else transforms
        if tuple(axes) in transforms:
            plan_fwd, plan_bck = transforms[tuple(axes)]
        else:
            if np.issubdtype(dtype, np.floating):
                plan_fwd = np.fft.rfftn
                plan_bck = np.fft.irfftn
            else:
                plan_fwd = np.fft.fftn
                plan_bck = np.fft.ifftn

        s = tuple(np.take(shape, axes))
        U = fftw.aligned(shape, dtype=dtype)
        V = plan_fwd(U, s=s, axes=axes).astype(dtype.char.upper()) # Numpy returns complex double if input single precision
        V = fftw.aligned_like(V)
        M = np.prod(s)

        # Numpy has forward transform unscaled and backward scaled with 1/N
        return (_Yfftn_wrap(plan_fwd, U, V, 1, {'s': s, 'axes': axes}),
                _Yfftn_wrap(plan_bck, V, U, M, {'s': s, 'axes': axes}))

    def _Xfftn_plan_mkl(shape, axes, dtype, transforms, options): #pragma: no cover

        transforms = {} if transforms is None else transforms
        if tuple(axes) in transforms:
            plan_fwd, plan_bck = transforms[tuple(axes)]
        else:
            if np.issubdtype(dtype, np.floating):
                from mkl_fft._numpy_fft import rfftn, irfftn
                plan_fwd = rfftn
                plan_bck = irfftn
            else:
                from mkl_fft._numpy_fft import fftn, ifftn
                plan_fwd = fftn
                plan_bck = ifftn

        s = tuple(np.take(shape, axes))
        U = fftw.aligned(shape, dtype=dtype)
        V = plan_fwd(U, s=s, axes=axes)
        V = fftw.aligned_like(V)
        M = np.prod(s)

        return (_Yfftn_wrap(plan_fwd, U, V, 1, {'s': s, 'axes': axes}),
                _Yfftn_wrap(plan_bck, V, U, M, {'s': s, 'axes': axes}))

    def _Xfftn_plan_scipy(shape, axes, dtype, transforms, options):

        transforms = {} if transforms is None else transforms
        if tuple(axes) in transforms:
            plan_fwd, plan_bck = transforms[tuple(axes)]
        else:
            from scipy.fftpack import fftn, ifftn # No rfftn/irfftn methods
            plan_fwd = fftn
            plan_bck = ifftn

        s = tuple(np.take(shape, axes))
        U = fftw.aligned(shape, dtype=dtype)
        V = plan_fwd(U, shape=s, axes=axes)
        V = fftw.aligned_like(V)
        M = np.prod(s)
        return (_Yfftn_wrap(plan_fwd, U, V, 1, {'shape': s, 'axes': axes}),
                _Yfftn_wrap(plan_bck, V, U, M, {'shape': s, 'axes': axes}))

    class _Yfftn_wrap(object):
        #Wraps numpy/scipy/mkl transforms to FFTW style
        # pylint: disable=too-few-public-methods

        __slots__ = ('_xfftn', '_M', '_opt', '__doc__', '_input_array', '_output_array')

        def __init__(self, xfftn_obj, input_array, output_array, M, opt):
            object.__setattr__(self, '_xfftn', xfftn_obj)
            object.__setattr__(self, '_opt', opt)
            object.__setattr__(self, '_M', M)
            object.__setattr__(self, '_input_array', input_array)
            object.__setattr__(self, '_output_array', output_array)
            object.__setattr__(self, '__doc__', xfftn_obj.__doc__)

        @property
        def input_array(self):
            return object.__getattribute__(self, '_input_array')

        @property
        def output_array(self):
            return object.__getattribute__(self, '_output_array')

        @property
        def xfftn(self):
            return object.__getattribute__(self, '_xfftn')

        @property
        def opt(self):
            return object.__getattribute__(self, '_opt')

        @property
        def M(self):
            return object.__getattribute__(self, '_M')

        def __call__(self, *args, **kwargs):
            self.opt.update(kwargs)
            self.output_array[...] = self.xfftn(self.input_array, **self.opt)
            if abs(self.M-1) > 1e-8:
                self._output_array *= self.M
            return self.output_array

    class _Xfftn_wrap(object):
        #Common interface for all serial transforms
        # pylint: disable=too-few-public-methods

        __slots__ = ('_xfftn', '__doc__', '_input_array', '_output_array')

        def __init__(self, xfftn_obj, input_array, output_array):
            object.__setattr__(self, '_xfftn', xfftn_obj)
            object.__setattr__(self, '_input_array', input_array)
            object.__setattr__(self, '_output_array', output_array)
            object.__setattr__(self, '__doc__', xfftn_obj.__doc__)

        @property
        def input_array(self):
            return object.__getattribute__(self, '_input_array')

        @property
        def output_array(self):
            return object.__getattribute__(self, '_output_array')

        @property
        def xfftn(self):
            return object.__getattribute__(self, '_xfftn')

        def __call__(self, input_array=None, output_array=None, **options):
            if input_array is not None:
                self.input_array[...] = input_array
            self.xfftn(**options)
            if output_array is not None:
                output_array[...] = self.output_array
                return output_array
            else:
                return self.output_array

    class FFTBase(object):
        """Base class for serial FFT transforms

        Parameters
        ----------
        shape : list or tuple of ints
            shape of input array planned for
        axes : None, int or tuple of ints, optional
            axes to transform over. If None transform over all axes
        dtype : np.dtype, optional
            Type of input array
        padding : bool, number or list of numbers
            If False, then no padding. If number, then apply this number as padding
            factor for all axes. If list of numbers, then each number gives the
            padding for each axis. Must be same length as axes.
        """

        def __init__(self, shape, axes=None, dtype=float, padding=False):
            shape = list(shape) if np.ndim(shape) else [shape]
            assert len(shape) > 0
            assert min(shape) > 0
            if axes is not None:
                axes = list(axes) if np.ndim(axes) else [axes]
                for i, axis in enumerate(axes):
                    if axis < 0:
                        axes[i] = axis + len(shape)
            else:
                axes = list(range(len(shape)))
            assert min(axes) >= 0
            assert max(axes) < len(shape)
            assert 0 < len(axes) <= len(shape)
            assert sorted(axes) == sorted(set(axes))

            dtype = np.dtype(dtype)
            assert dtype.char in 'fdgFDG'
            self.shape = shape
            self.axes = axes
            self.dtype = dtype
            self.padding = padding
            self.real_transform = np.issubdtype(dtype, np.floating)
            self.padding_factor = 1

        def _truncation_forward(self, padded_array, trunc_array):
            axis = self.axes[-1]
            if self.padding_factor > 1.0+1e-8:
                trunc_array.fill(0)
                N0 = self.forward.output_array.shape[axis]
                if self.real_transform:
                    N = trunc_array.shape[axis]
                    s = [slice(None)]*trunc_array.ndim
                    s[axis] = slice(0, N)
                    trunc_array[:] = padded_array[tuple(s)]
                    if N0 % 2 == 0:
                        s[axis] = N-1
                        s = tuple(s)
                        trunc_array[s] = trunc_array[s].real
                        trunc_array[s] *= 2
                else:
                    N = trunc_array.shape[axis]
                    su = [slice(None)]*trunc_array.ndim
                    su[axis] = slice(0, N//2+1)
                    trunc_array[tuple(su)] = padded_array[tuple(su)]
                    su[axis] = slice(-(N//2), None)
                    trunc_array[tuple(su)] += padded_array[tuple(su)]

        def _padding_backward(self, trunc_array, padded_array):
            axis = self.axes[-1]
            if self.padding_factor > 1.0+1e-8:
                padded_array.fill(0)
                N0 = self.forward.output_array.shape[axis]
                if self.real_transform:
                    s = [slice(0, n) for n in trunc_array.shape]
                    padded_array[tuple(s)] = trunc_array[:]
                    N = trunc_array.shape[axis]
                    if N0 % 2 == 0: # Symmetric Fourier interpolator
                        s[axis] = N-1
                        s = tuple(s)
                        padded_array[s] = padded_array[s].real
                        padded_array[s] *= 0.5
                else:
                    N = trunc_array.shape[axis]
                    su = [slice(None)]*trunc_array.ndim
                    su[axis] = slice(0, N//2+1)
                    padded_array[tuple(su)] = trunc_array[tuple(su)]
                    su[axis] = slice(-(N//2), None)
                    padded_array[tuple(su)] = trunc_array[tuple(su)]
                    if N0 % 2 == 0:  # Use symmetric Fourier interpolator
                        su[axis] = N//2
                        padded_array[tuple(su)] *= 0.5
                        su[axis] = -(N//2)
                        padded_array[tuple(su)] *= 0.5

    class FFT(FFTBase):
        """Class for serial FFT transforms

        Parameters
        ----------
        shape : list or tuple of ints
            shape of input array planned for
        axes : None, int or tuple of ints, optional
            axes to transform over. If None transform over all axes
        dtype : np.dtype, optional
            Type of input array
        padding : bool, number or list of numbers
            If False, then no padding. If number, then apply this number as padding
            factor for all axes. If list of numbers, then each number gives the
            padding for each axis. Must be same length as axes.
        backend : str, optional
            Choose backend for serial transforms (``fftw``, ``pyfftw``, ``numpy``,
            ``scipy``, ``mkl_fft``). Default is ``fftw``
        transforms : None or dict, optional
            Dictionary of axes to serial transforms (forward and backward) along
            those axes. For example::

                {(0, 1): (dctn, idctn), (2, 3): (dstn, idstn)}

            If missing the default is to use rfftn/irfftn for real input arrays and
            fftn/ifftn for complex input arrays. Real-to-real transforms can be
            configured using this dictionary and real-to-real transforms from the
            :mod:`.fftw.xfftn` module.
        kw : dict
            Parameters passed to serial transform object

        Methods
        -------
        forward(input_array=None, output_array=None, **options)
            Generic serial forward transform

            Parameters
            ----------
            input_array : array, optional
            output_array : array, optional
            options : dict
                parameters to serial transforms

            Returns
            -------
            output_array : array

        backward(input_array=None, output_array=None, **options)
            Generic serial backward transform

            Parameters
            ----------
            input_array : array, optional
            output_array : array, optional
            options : dict
                parameters to serial transforms

            Returns
            -------
            output_array : array

        """
        def __init__(self, shape, axes=None, dtype=float, padding=False,
                     backend='fftw', transforms=None, **kw):
            FFTBase.__init__(self, shape, axes, dtype, padding)
            plan = {
                'pyfftw': _Xfftn_plan_pyfftw,
                'fftw': _Xfftn_plan_fftw,
                'numpy': _Xfftn_plan_numpy,
                'mkl_fft': _Xfftn_plan_mkl,
                'scipy': _Xfftn_plan_scipy,
            }[backend]
            self.backend = backend
            self.fwd, self.bck = plan(self.shape, self.axes, self.dtype, transforms, kw)
            U, V = self.fwd.input_array, self.fwd.output_array
            self.M = 1
            if backend != 'fftw':
                self.M = 1./np.prod(np.take(self.shape, self.axes))
            else:
                self.M = self.fwd.get_normalization()
            if backend == 'scipy':
                self.real_transform = False # No rfftn/irfftn methods
            self.padding_factor = 1.0
            if padding is not False:
                self.padding_factor = padding[self.axes[-1]] if np.ndim(padding) else padding
            if abs(self.padding_factor-1.0) > 1e-8:
                assert len(self.axes) == 1
                trunc_array = self._get_truncarray(shape, V.dtype)
                self.forward = _Xfftn_wrap(self._forward, U, trunc_array)
                self.backward = _Xfftn_wrap(self._backward, trunc_array, U)
            else:
                self.forward = _Xfftn_wrap(self._forward, U, V)
                self.backward = _Xfftn_wrap(self._backward, V, U)

        def _forward(self, **kw):
            normalize = kw.pop('normalize', True)
            self.fwd(None, None, **kw)
            self._truncation_forward(self.fwd.output_array, self.forward.output_array)
            if normalize:
                self.forward._output_array *= self.M
            return self.forward.output_array

        def _backward(self, **kw):
            normalize = kw.pop('normalize', False)
            self._padding_backward(self.backward.input_array, self.bck.input_array)
            self.bck(None, None, **kw)
            if normalize:
                self.backward._output_array *= self.M
            return self.backward.output_array

        def _get_truncarray(self, shape, dtype):
            axis = self.axes[-1]
            if not self.real_transform:
                shape = list(shape)
                shape[axis] = int(np.round(shape[axis] / self.padding_factor))
                return fftw.aligned(shape, dtype=dtype)

            shape = list(shape)
            shape[axis] = int(np.round(shape[axis] / self.padding_factor))
            shape[axis] = shape[axis]//2 + 1
            return fftw.aligned(shape, dtype=dtype)