import os
import numpy
import logging

from nbodykit import CurrentMPIComm
from nbodykit.binned_statistic import BinnedStatistic
from nbodykit.meshtools import SlabIterator
from nbodykit.base.catalog import CatalogSourceBase
from nbodykit.base.mesh import MeshSource

from .fftpower import FFTBase
from .fftpower import project_to_basis
from .fftpower import _find_unique_edges

class FFTCorr(FFTBase):
    r"""
    Algorithm to compute the 1d or 2d correlation and/or multipoles
    in a periodic box, using a Fast Fourier Transform (FFT).

    This computes the power spectrum as the square of the Fourier modes of the
    density field, which are computed via a FFT. Then it is transformed back
    to obtain the correlation function.

    Results are computed when the object is inititalized. See the documenation
    of :func:`~FFTCorr.run` for the attributes storing the results.

    .. note::
        This is very similar to :class:`~nbodykit.algorithms.fftpower.FFTPower`. 

    Parameters
    ----------
    first : CatalogSource, MeshSource
        the source for the first field; if a CatalogSource is provided, it
        is automatically converted to MeshSource using the default painting
        parameters (via :func:`~nbodykit.base.catalogmesh.CatalogMesh.to_mesh`)
    mode : {'1d', '2d'}
        compute either 1d or 2d power spectra
    Nmesh : int, optional
        the number of cells per side in the particle mesh used to paint the source
    BoxSize : int, 3-vector, optional
        the size of the box
    second : CatalogSource, MeshSource, optional
        the second source for cross-correlations
    los : array_like , optional
        the direction to use as the line-of-sight; must be a unit vector
    Nmu : int, optional
        the number of mu bins to use from :math:`\mu=[0,1]`;
        if `mode = 1d`, then ``Nmu`` is set to 1
    dr : float, optional
        the linear spacing of ``r`` bins to use; if not provided, the
        fundamental mode  of the box is used; if `dr=0`, the bins are tight, such
        that each bin has a unique r value.
    rmin : float, optional
        the lower edge of the first ``r`` bin to use
    rmax : float, optional
        the upper limit of the last ``r`` bin to use
    poles : list of int, optional
        a list of multipole numbers ``ell`` to compute :math:`\xi_\ell(r)`
        from :math:`\xi(r,\mu)`
    """
    logger = logging.getLogger('FFTCorr')

    def __init__(self, first, mode, Nmesh=None, BoxSize=None, second=None,
                    los=[0, 0, 1], Nmu=5, dr=None, rmin=0., rmax=None, poles=[]):

        # mode is either '1d' or '2d'
        if mode not in ['1d', '2d']:
            raise ValueError("`mode` should be either '1d' or '2d'")

        if poles is None:
            poles = []

        # check los
        if numpy.isscalar(los) or len(los) != 3:
            raise ValueError("line-of-sight ``los`` should be vector with length 3")
        if not numpy.allclose(numpy.einsum('i,i', los, los), 1.0, rtol=1e-5):
            raise ValueError("line-of-sight ``los`` must be a unit vector")

        FFTBase.__init__(self, first, second, Nmesh, BoxSize)

        # save meta-data
        self.attrs['mode'] = mode
        self.attrs['los'] = los
        self.attrs['Nmu'] = Nmu
        self.attrs['poles'] = poles

        if dr is None:
            dr = self.attrs['BoxSize'].min() / self.attrs['Nmesh'].max()

        self.attrs['dr'] = dr
        self.attrs['rmin'] = rmin
        self.attrs['rmax'] = rmax

        self.corr, self.poles = self.run()

        # compatability
        self.attrs.update(self.corr.attrs)

    def run(self):
        r"""
        Compute the correlation function in a periodic box, using FFTs.

        returns
        -------
        corr : :class:`~nbodykit.binned_statistic.BinnedStatistic`
            a BinnedStatistic object that holds the measured :math:`\xi(r)` or
            :math:`\xi(r,\mu)`. It stores the following variables:

            - r :
                the mean value for each ``r`` bin
            - mu : ``mode=2d`` only
                the mean value for each ``mu`` bin
            - corr :
                real array storing the correlation function
            - modes :
                the number of modes averaged together in each bin

        poles : :class:`~nbodykit.binned_statistic.BinnedStatistic` or ``None``
            a BinnedStatistic object to hold the multipole results
            :math:`\xi_\ell(r)`; if no multipoles were requested by the user,
            this is ``None``. It stores the following variables:

            - r :
                the mean value for each ``r`` bin
            - power_L :
                complex array storing the real and imaginary components for
                the :math:`\ell=L` multipole
            - modes :
                the number of modes averaged together in each bin

        corr.attrs, poles.attrs : dict
            dictionary of meta-data; in addition to storing the input parameters,
            it includes the following fields computed during the algorithm
            execution:

            - shotnoise : float
                the power Poisson shot noise, equal to :math:`V/N`, where
                :math:`V` is the volume of the box and `N` is the total
                number of objects; if a cross-correlation is computed, this
                will be equal to zero
            - N1 : int
                the total number of objects in the first source
            - N2 : int
                the total number of objects in the second source
        """

        # only need one mu bin if 1d case is requested
        if self.attrs['mode'] == "1d": self.attrs['Nmu'] = 1

        # measure the 3D power (y3d is a ComplexField)
        y3d, attrs = self._compute_3d_power(self.first, self.second)

        # measure the 3D correlation (y3d is a RealField)
        y3d = y3d.c2r(out=Ellipsis)

        # correlation is dimensionless
        # Note that L^3 cancels with dk^3.
        y3d[...] *= 1.0 / y3d.BoxSize.prod()

        # binning in k out to the minimum nyquist frequency
        # (accounting for possibly anisotropic box)
        dr = self.attrs['dr']
        rmin = self.attrs['rmin']
        rmax = self.attrs['rmax']
        if rmax is None:
            rmax = 0.5 * y3d.BoxSize.min() + dr/2
        if dr > 0:
            redges = numpy.arange(rmin, rmax, dr)
            rcenters = None
        else:
            redges, rcenters = _find_unique_edges(y3d.x, y3d.BoxSize / y3d.Nmesh, rmax, self.comm)


        # project on to the desired basis
        muedges = numpy.linspace(0, 1, self.attrs['Nmu']+1, endpoint=True)
        edges = [redges, muedges]
        coords = [rcenters, None]
        result, pole_result = project_to_basis(y3d, edges,
                                               poles=self.attrs['poles'],
                                               los=self.attrs['los'])

        # format the corr results into structured array
        if self.attrs['mode'] == "1d":
            cols = ['r', 'corr', 'modes']
            icols = [0, 2, 3]
            edges = edges[0:1]
            coords = coords[0:1]
        else:
            cols = ['r', 'mu', 'corr', 'modes']
            icols = [0, 1, 2, 3]

        # corr results as a structured array
        dtype = numpy.dtype([(name, result[icol].dtype.str) for icol,name in zip(icols,cols)])
        corr = numpy.squeeze(numpy.empty(result[0].shape, dtype=dtype))
        for icol, col in zip(icols, cols):
            corr[col][:] = numpy.squeeze(result[icol])

        # multipole results as a structured array
        poles = None
        if pole_result is not None:
            r, poles, N = pole_result
            cols = ['r'] + ['corr_%d' %l for l in self.attrs['poles']] + ['modes']
            result = [r] + [pole for pole in poles] + [N]

            dtype = numpy.dtype([(name, result[icol].dtype.str) for icol,name in enumerate(cols)])
            poles = numpy.empty(result[0].shape, dtype=dtype)
            for icol, col in enumerate(cols):
                poles[col][:] = result[icol]

        # set all the necessary results
        return self._make_datasets(edges, poles, corr, coords, attrs)

    def __getstate__(self):
        state = dict(
                    corr=self.corr.__getstate__(),
                    poles=self.poles.__getstate__() if self.poles is not None else None,
                    attrs=self.attrs)
        return state

    def __setstate__(self, state):
        self.attrs = state['attrs']
        self.corr = BinnedStatistic.from_state(state['corr'])
        if state['poles'] is not None:
            self.poles = BinnedStatistic.from_state(state['poles'])

    def _make_datasets(self, edges, poles, corr, coords, attrs):

        if self.attrs['mode'] == '1d':
            corr = BinnedStatistic(['r'], edges, corr, fields_to_sum=['modes'], coords=coords, **attrs)
        else:
            corr = BinnedStatistic(['r', 'mu'], edges, corr, fields_to_sum=['modes'], coords=coords, **attrs)
        if poles is not None:
            poles = BinnedStatistic(['r'], [corr.edges['r']], poles, fields_to_sum=['modes'], coords=coords, **attrs)

        return corr, poles
