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)