# Source code for mpi4py_fft.mpifft

    import numpy as np

    from .libfft import FFT
    from .pencil import Pencil
    from .pencil import Subcomm

    class Transform(object):
        """Class for performing any parallel transform, forward or backward

        Parameters
        ----------
        xfftn : list of serial transform objects
        transfer : list of global redistribution objects
        pencil : list of two pencil objects
            The two pencils represent the input and final output configuration of
            the distributed global arrays

        """
        def __init__(self, xfftn, transfer, pencil):
            assert len(xfftn) == len(transfer) + 1 and len(pencil) == 2
            self._xfftn = tuple(xfftn)
            self._transfer = tuple(transfer)
            self._pencil = tuple(pencil)

        @property
        def input_array(self):
            """Return input array of Transform"""
            return self._xfftn[0].input_array

        @property
        def output_array(self):
            """Return output array of Transform"""
            return self._xfftn[-1].output_array

        @property
        def input_pencil(self):
            """Return input pencil of Transform"""
            return self._pencil[0]

        @property
        def output_pencil(self):
            """Return output pencil of Transform"""
            return self._pencil[1]

        def __call__(self, input_array=None, output_array=None, **kw):
            """Compute transform

            Parameters
            ----------
            input_array : array, optional
            output_array : array, optional
            kw : dict
                parameters to serial transforms
                Note in particular that the keyword 'normalize'=True/False can be
                used to turn normalization on or off. Default is to enable
                normalization for forward transforms and disable it for backward.

            Note
            ----
            If input_array/output_array are not given, then use predefined arrays
            as planned with serial transform object _xfftn.

            """
            if input_array is not None:
                self.input_array[...] = input_array

            for i in range(len(self._transfer)):
                self._xfftn[i](**kw)
                arrayA = self._xfftn[i].output_array
                arrayB = self._xfftn[i+1].input_array
                self._transfer[i](arrayA, arrayB)
            self._xfftn[-1](**kw)

            if output_array is not None:
                output_array[...] = self.output_array
                return output_array
            else:
                return self.output_array

    class PFFT(object):
        """Base class for parallel FFT transforms

        Parameters
        ----------
        comm : MPI communicator
        shape : sequence of ints, optional
            shape of input array planned for
        axes : None, int, sequence of ints or sequence of sequence of ints, optional
            axes to transform over.

            - None -> All axes are transformed
            - int -> Just one axis to transform over
            - sequence of ints -> e.g., (0, 1, 2) or (0, 2, 1)
            - sequence of sequence of ints -> e.g., ((0,), (1,)) or ((0,), (1, 2))
              For seq. of seq. of ints all but the last transformed sequence
              may be longer than 1. This corresponds to collapsing axes, where
              serial FFTs are performed for all collapsed axes in one single call
        dtype : np.dtype, optional
            Type of input array
        grid : sequence of ints, optional
            Define processor grid sizes. Non positive values act as wildcards to
            allow MPI compute optimal decompositions. The sequence is padded with
            ones to match the global transform dimension.
            Use ``(-1,)`` to get a slab decomposition on the first axis.
            Use ``(1, -1)`` to get a slab decomposition  on the second axis.
            Use ``(P, Q)`` or ``(P, Q, 1)`` to get a 3D transform with 2D-pencil
            decomposition on a PxQ processor grid with the last axis non distributed.
            Use ``(P, 1, Q)`` to get a 3D transform with 2D-pencil decomposition on
            a PxQ processor grid with the second to last axis non distributed.
        padding : bool, number or sequence of numbers, optional
            If False, then no padding. If number, then apply this number as padding
            factor for all axes. If sequence of numbers, then each number gives the
            padding for each axis. Must be same length as axes.
        collapse : bool, optional
            If True try to collapse several serial transforms into one
        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. See Examples.

        Other Parameters
        ----------------
        darray : DistArray object, optional
            Create PFFT using information contained in ``darray``, neglecting most
            optional Parameters above
        slab : bool or int, optional
            DEPRECATED. If True then distribute only one axis of the global array.

        Methods
        -------
        forward(input_array=None, output_array=None, **kw)
            Parallel forward transform. The method is an instance of the
            :class:`.Transform` class. See :meth:`.Transform.__call__`

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

            Returns
            -------
            output_array : array

        backward(input_array=None, output_array=None, **kw)
            Parallel backward transform. The method is an instance of the
            :class:`.Transform` class. See :meth:`.Transform.__call__`

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

            Returns
            -------
            output_array : array

        Examples
        --------
        >>> import numpy as np
        >>> from mpi4py import MPI
        >>> from mpi4py_fft import PFFT, newDistArray
        >>> N = np.array([12, 14, 15], dtype=int)
        >>> fft = PFFT(MPI.COMM_WORLD, N, axes=(0, 1, 2))
        >>> u = newDistArray(fft, False)
        >>> u[:] = np.random.random(u.shape).astype(u.dtype)
        >>> u_hat = fft.forward(u)
        >>> uj = np.zeros_like(u)
        >>> uj = fft.backward(u_hat, uj)
        >>> assert np.allclose(uj, u)

        Now configure with real-to-real discrete cosine transform type 3

        >>> from mpi4py_fft.fftw import rfftn, irfftn, dctn, idctn
        >>> import functools
        >>> dct = functools.partial(dctn, type=3)
        >>> idct = functools.partial(idctn, type=3)
        >>> transforms = {(1, 2): (dct, idct)}
        >>> r2c = PFFT(MPI.COMM_WORLD, N, axes=((0,), (1, 2)), transforms=transforms)
        >>> u = newDistArray(r2c, False)
        >>> u[:] = np.random.random(u.shape).astype(u.dtype)
        >>> u_hat = r2c.forward(u)
        >>> uj = np.zeros_like(u)
        >>> uj = r2c.backward(u_hat, uj)
        >>> assert np.allclose(uj, u)

        """
        def __init__(self, comm, shape=None, axes=None, dtype=float, grid=None,
                     padding=False, collapse=False, backend='fftw',
                     transforms=None, darray=None, **kw):
            # pylint: disable=too-many-locals
            # pylint: disable=too-many-branches
            # pylint: disable=too-many-statements

            if shape is None:
                assert darray is not None
                shape = darray.pencil.shape

            if axes is not None:
                axes = list(axes) if not isinstance(axes, int) else [axes]
            else:
                axes = list(range(len(shape)))
                if darray is not None:
                    # Make sure aligned axis of darray is transformed first
                    axes = list(np.roll(axes, len(shape)-1-darray.alignment))

            for i, ax in enumerate(axes):
                if isinstance(ax, (int, np.integer)):
                    if ax < 0:
                        ax += len(shape)
                    axes[i] = (ax,)
                else:
                    assert isinstance(ax, (tuple, list))
                    ax = list(ax)
                    for j, a in enumerate(ax):
                        assert isinstance(a, int)
                        if a < 0:
                            a += len(shape)
                            ax[j] = a
                    axes[i] = ax
                assert min(axes[i]) >= 0
                assert max(axes[i]) < len(shape)
                assert 0 < len(axes[i]) <= len(shape)
                assert sorted(axes[i]) == sorted(set(axes[i]))

            self.axes = axes
            shape = list(shape)

            if darray is None:
                dtype = np.dtype(dtype)
                assert dtype.char in 'fdgFDG'

                if padding is not False:
                    assert len(padding) == len(shape)
                    for ax in axes:
                        if len(ax) == 1 and padding[ax[0]] > 1.0+1e-6:
                            old = float(shape[ax[0]])
                            shape[ax[0]] = int(np.floor(shape[ax[0]]*padding[ax[0]]))
                            padding[ax[0]] = shape[ax[0]] / old

                self._input_shape = tuple(shape)
                assert len(shape) > 0
                assert min(shape) > 0

                slab = kw.pop('slab', False)

                if grid is not None:
                    assert not isinstance(comm, Subcomm)
                    assert slab is False
                    grid = tuple(grid)
                    assert len(grid) <= len(shape)
                    dims = list(grid) + [1] * (len(shape) - len(grid))
                    comm = Subcomm(comm, dims)

                if isinstance(comm, Subcomm):
                    assert slab is False
                    assert len(comm) == len(shape)
                    assert np.all([comm[ax].Get_size() == 1 for ax in axes[-1]])
                    self.subcomm = comm
                else:
                    if slab is False or slab is None:
                        dims = [0] * len(shape)
                        for ax in axes[-1]:
                            dims[ax] = 1
                    else: #pragma: no cover
                        if slab is True:
                            axis = (axes[-1][-1] + 1) % len(shape)
                        else:
                            axis = slab
                            if axis < 0:
                                axis = axis + len(shape)
                            assert 0 <= axis < len(shape)
                        dims = [1] * len(shape)
                        dims[axis] = comm.Get_size()

                    self.subcomm = Subcomm(comm, dims)
            else:
                dtype = darray.dtype
                self.subcomm = darray.subcomm
                self._input_shape = tuple(shape)
                commsizes = darray.commsizes
                assert np.all([commsizes[ax] == 1 for ax in axes[-1]]), "Set keyword axes such that axes to transform first are aligned"

            self.collapse = collapse
            if collapse is True:
                groups = [[]]
                for ax in reversed(axes):
                    if np.all([self.subcomm[axis].Get_size() == 1 for axis in ax]):
                        [groups[0].insert(0, axis) for axis in reversed(ax)]
                    else:
                        groups.insert(0, ax)
                axes = groups

            self.axes = tuple(map(tuple, axes))
            self.xfftn = []
            self.transfer = []
            self.pencil = [None, None]

            axes = self.axes[-1]
            pencil = Pencil(self.subcomm, shape, axes[-1])
            xfftn = FFT(pencil.subshape, axes, dtype, padding, backend=backend,
                        transforms=transforms, **kw)
            self.xfftn.append(xfftn)
            self.pencil[0] = pencilA = pencil
            if not shape[axes[-1]] == xfftn.forward.output_array.shape[axes[-1]]:
                dtype = xfftn.forward.output_array.dtype
                shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
                pencilA = Pencil(self.subcomm, shape, axes[-1])

            for axes in reversed(self.axes[:-1]):
                pencilB = pencilA.pencil(axes[-1])
                transAB = pencilA.transfer(pencilB, dtype)
                xfftn = FFT(pencilB.subshape, axes, dtype, padding, backend=backend,
                            transforms=transforms, **kw)
                self.xfftn.append(xfftn)
                self.transfer.append(transAB)
                pencilA = pencilB
                if not shape[axes[-1]] == xfftn.forward.output_array.shape[axes[-1]]:
                    dtype = xfftn.forward.output_array.dtype
                    shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
                    pencilA = Pencil(pencilB.subcomm, shape, axes[-1])

            self.pencil[1] = pencilA
            self._output_shape = tuple(shape)

            self.forward = Transform(
                [o.forward for o in self.xfftn],
                [o.forward for o in self.transfer],
                self.pencil)
            self.backward = Transform(
                [o.backward for o in self.xfftn[::-1]],
                [o.backward for o in self.transfer[::-1]],
                self.pencil[::-1])

        def destroy(self):
            if isinstance(self.subcomm, Subcomm):
                self.subcomm.destroy()
            for trans in self.transfer:
                trans.destroy()

        def shape(self, forward_output=True):
            """The local (to each processor) shape of data

            Parameters
            ----------
            forward_output : bool, optional
                Return shape of output array (spectral space) if True, else return
                shape of input array (physical space)
            """
            if forward_output is not True:
                return self.forward.input_pencil.subshape
            return self.forward.output_array.shape

        def local_slice(self, forward_output=True):
            """The local view into the global data

            Parameters
            ----------
            forward_output : bool, optional
                Return local slices of output array (spectral space) if True, else
                return local slices of input array (physical space)

            """
            if forward_output is not True:
                ip = self.forward.input_pencil
                s = [slice(start, start+shape) for start, shape in zip(ip.substart,
                                                                       ip.subshape)]
            else:
                ip = self.backward.input_pencil
                s = [slice(start, start+shape) for start, shape in zip(ip.substart,
                                                                       ip.subshape)]
            return tuple(s)

        def global_shape(self, forward_output=False):
            """Return global shape of associated tensors

            Parameters
            ----------
            forward_output : bool, optional
                If True then return global shape of spectral space, i.e., the input
                to a backward transfer. If False then return shape of physical
                space, i.e., the input to a forward transfer.
            """
            if forward_output:
                return self._output_shape
            return self._input_shape

        @property
        def dimensions(self):
            """The number of dimensions for transformed arrays"""
            return len(self.forward.input_array.shape)

        def dtype(self, forward_output=False):
            """The type of transformed arrays

            Parameters
            ----------
            forward_output : bool, optional
                If True then return dtype of an array that is the result of a
                forward transform. Otherwise, return the dtype of an array that
                is input to a forward transform.
            """
            if forward_output:
                return self.forward.output_array.dtype
            return self.forward.input_array.dtype