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 nbodykit.base.catalogmesh import CatalogMesh

class FFTBase(object):
    """
    Base class provides functions for periodic FFT based Power spectrum code.

    Parameters
    ----------
    first : CatalogSource
        the first catalog source
    second : CatalogSource, None
        the second source, or None for auto-correlations
    Nmesh : int, 3-vector
        the number of cells per mesh size
    BoxSize : 3-vector
        the size of the box
    """
    def __init__(self, first, second, Nmesh, BoxSize):
        from pmesh.pm import ParticleMesh

        first = _cast_source(first, Nmesh=Nmesh, BoxSize=BoxSize)
        if second is not None:
            second = _cast_source(second, Nmesh=Nmesh, BoxSize=BoxSize)
        else:
            second = first

        self.first = first
        self.second = second

        # grab comm from first source
        self.comm = first.comm

        # check for comm mismatch
        assert second.comm is first.comm, "communicator mismatch between input sources"

        # check box size
        if not numpy.array_equal(first.attrs['BoxSize'], second.attrs['BoxSize']):
            raise ValueError("'BoxSize' mismatch between sources in FFTPower")

        # save meta-data
        self.attrs = {}
        self.attrs['Nmesh'] = first.attrs['Nmesh'].copy()
        self.attrs['BoxSize'] = first.attrs['BoxSize'].copy()

        self.attrs.update(zip(['Lx', 'Ly', 'Lz'], self.attrs['BoxSize']))
        self.attrs.update({'volume':self.attrs['BoxSize'].prod()})

    def save(self, output):
        """
        Save the result to disk. The format is currently JSON.
        """
        import json
        from nbodykit.utils import JSONEncoder

        # only the master rank writes
        if self.comm.rank == 0:
            self.logger.info('measurement done; saving result to %s' % output)

            with open(output, 'w') as ff:
                json.dump(self.__getstate__(), ff, cls=JSONEncoder)

    @classmethod
    @CurrentMPIComm.enable
    def load(cls, output, comm=None):
        """
        Load a saved result. The result has been saved to disk with :func:`save`.
        """
        import json
        from nbodykit.utils import JSONDecoder
        if comm.rank == 0:
            with open(output, 'r') as ff:
                state = json.load(ff, cls=JSONDecoder)
        else:
            state = None
        state = comm.bcast(state)
        self = object.__new__(cls)
        self.__setstate__(state)
        self.comm = comm

        return self

    def _compute_3d_power(self):
        """
        Compute and return the 3D power from two input sources

        Returns
        -------
        p3d : array_like (complex)
            the 3D complex array holding the power spectrum
        """
        c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])

        # compute the auto power of single supplied field
        if self.first is self.second:
            c2 = c1
        else:
            c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])

        # calculate the 3d power spectrum, slab-by-slab to save memory
        p3d = c1
        for (s0, s1, s2) in zip(p3d.slabs, c1.slabs, c2.slabs):
            s0[...] = s1 * s2.conj()

        for i, s0 in zip(p3d.slabs.i, p3d.slabs):
            # clear the zero mode.
            mask = True
            for i1 in i:
                mask = mask & (i1 == 0)
            s0[mask] = 0

        # the complex field is dimensionless; power is L^3
        # ref to http://icc.dur.ac.uk/~tt/Lectures/UA/L4/cosmology.pdf
        p3d[...] *= self.attrs['BoxSize'].prod()

        # get the number of objects (in a safe manner)
        N1 = c1.attrs.get('N', 0)
        N2 = c2.attrs.get('N', 0)
        self.attrs.update({'N1':N1, 'N2':N2})

        # add shotnoise (nonzero only for auto-spectra)
        Pshot = 0
        if self.first is self.second:
            if 'shotnoise' not in c1.attrs:
                if isinstance(self.first, CatalogMesh):
                    import warnings
                    warnings.warn(("no 'shotnoise' found for auto power spectrum "
                                   "of discrete data in FFTPower"))
            else:
                Pshot = c1.attrs['shotnoise']
        self.attrs['shotnoise'] = Pshot

        return p3d

class FFTPower(FFTBase):
    """
    Algorithm to compute the 1d or 2d power spectrum 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.

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

    .. note::
        A full tutorial on the class is available in the documentation
        :ref:`here <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
    dk : float, optional
        the linear spacing of ``k`` bins to use; if not provided, the
        fundamental mode  of the box is used
    kmin : float, optional
        the lower edge of the first ``k`` bin to use
    poles : list of int, optional
        a list of multipole numbers ``ell`` to compute :math:`P_\ell(k)`
        from :math:`P(k,\mu)`
    """
    logger = logging.getLogger('FFTPower')

    def __init__(self, first, mode, Nmesh=None, BoxSize=None, second=None,
                    los=[0, 0, 1], Nmu=5, dk=None, kmin=0., 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 dk is None:
            dk = 2 * numpy.pi / self.attrs['BoxSize'].min()

        self.attrs['dk'] = dk
        self.attrs['kmin'] = kmin

        self.run()

    def run(self):
        """
        Compute the power spectrum in a periodic box, using FFTs. This
        function returns nothing, but attaches several attributes
        to the class:

        - :attr:`edges`
        - :attr:`power`
        - :attr:`poles`

        Attributes
        ----------
        edges : array_like
            the edges of the wavenumber bins
        power : :class:`~nbodykit.binned_statistic.BinnedStatistic`
            a BinnedStatistic object that holds the measured :math:`P(k)` or
            :math:`P(k,\mu)`. It stores the following variables:

            - k :
                the mean value for each ``k`` bin
            - mu : ``mode=2d`` only
                the mean value for each ``mu`` bin
            - power :
                complex array storing the real and imaginary components of the power
            - modes :
                the number of Fourier modes averaged together in each bin

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

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

        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 = self._compute_3d_power()

        # binning in k out to the minimum nyquist frequency
        # (accounting for possibly anisotropic box)
        dk = self.attrs['dk']
        kmin = self.attrs['kmin']
        kedges = numpy.arange(kmin, numpy.pi*y3d.Nmesh.min()/y3d.BoxSize.max() + dk/2, dk)

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

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

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

        # multipole results as a structured array
        poles = None
        if pole_result is not None:
            k, poles, N = pole_result
            cols = ['k'] + ['power_%d' %l for l in self.attrs['poles']] + ['modes']
            result = [k] + [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
        self.edges = edges
        self.poles = poles
        self.power = power

        self._make_datasets()

    def __getstate__(self):
        state = dict(
                     edges=self.edges,
                     power=self.power.data,
                     poles=getattr(self.poles, 'data', None),
                     attrs=self.attrs)
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._make_datasets()

    def _make_datasets(self):

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

    def _compute_3d_power(self):
        """
        Compute and return the 3D power from two input sources

        Returns
        -------
        p3d : array_like (complex)
            the 3D complex array holding the power spectrum
        """
        c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])

        # compute the auto power of single supplied field
        if self.first is self.second:
            c2 = c1
        else:
            c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])

        # calculate the 3d power spectrum, slab-by-slab to save memory
        p3d = c1
        for (s0, s1, s2) in zip(p3d.slabs, c1.slabs, c2.slabs):
            s0[...] = s1 * s2.conj()

        for i, s0 in zip(p3d.slabs.i, p3d.slabs):
            # clear the zero mode.
            mask = True
            for i1 in i:
                mask = mask & (i1 == 0)
            s0[mask] = 0

        # the complex field is dimensionless; power is L^3
        # ref to http://icc.dur.ac.uk/~tt/Lectures/UA/L4/cosmology.pdf
        p3d[...] *= self.attrs['BoxSize'].prod()

        # get the number of objects (in a safe manner)
        N1 = c1.attrs.get('N', 0)
        N2 = c2.attrs.get('N', 0)
        self.attrs.update({'N1':N1, 'N2':N2})

        # add shotnoise (nonzero only for auto-spectra)
        Pshot = 0
        if self.first is self.second:
            if 'shotnoise' not in c1.attrs:
                if isinstance(self.first, CatalogMesh):
                    import warnings
                    warnings.warn(("no 'shotnoise' found for auto power spectrum "
                                   "of discrete data in FFTPower"))
            else:
                Pshot = c1.attrs['shotnoise']
        self.attrs['shotnoise'] = Pshot

        return p3d

class ProjectedFFTPower(FFTBase):
    """
    The power spectrum of a field in a periodic box, projected over certain axes.

    This is not really always physically meaningful, but convenient for
    making sense of Lyman-Alpha forest or lensing maps.

    This is usually called the 1d power spectrum or 2d power spectrum.

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

    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`)
    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
    axes : tuple
        axes to measure the power on. The axes not in the list will be averaged out.
        For example:
        - (0, 1) : project to x,y and measure power
        - (0) : project to x and measure power.
    dk : float, optional
        the linear spacing of ``k`` bins to use; if not provided, the
        fundamental mode  of the box is used
    kmin : float, optional
        the lower edge of the first ``k`` bin to use
    """
    logger = logging.getLogger('ProjectedFFTPower')

    def __init__(self, first, Nmesh=None, BoxSize=None, second=None,
                    axes=(0, 1), dk=None, kmin=0.):

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

        # only deal with 1d and 2d projections.
        assert len(axes) in (1, 2), "length of ``axes`` in ProjectedFFTPower should be 1 or 2"

        if dk is None:
            dk = 2 * numpy.pi / self.attrs['BoxSize'].min()

        self.attrs['dk'] = dk
        self.attrs['kmin'] = kmin

        self.attrs['axes'] = axes
        self.run()

    def run(self):
        """
        Run the algorithm. This attaches the following attributes to the class:

        - :attr:`edges`
        - :attr:`power`

        Attributes
        ----------
        edges : array_like
            the edges of the wavenumber bins
        power : :class:`~nbodykit.binned_statistic.BinnedStatistic`
            a BinnedStatistic object that holds the projected power.
            It stores the following variables:

            - k :
                the mean value for each ``k`` bin
            - power :
                complex array holding the real and imaginary components of the
                projected power
            - modes :
                the number of Fourier modes averaged together in each bin
        """
        c1 = self.first.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
        r1 = c1.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
        # average along projected axes;
        # part of product is the rfftn vs r2c (for axes)
        # the rest is for the mean (Nmesh - axes)
        c1 = numpy.fft.rfftn(r1) / self.attrs['Nmesh'].prod()

        # compute the auto power of single supplied field
        if self.first is self.second:
            c2 = c1
        else:
            c2 = self.second.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
            r2 = c2.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
            c2 = numpy.fft.rfftn(r2) / self.attrs['Nmesh'].prod() # average along projected axes

        pk = c1 * c2.conj()
        # clear the zero mode
        pk.flat[0] = 0

        shape = numpy.array([self.attrs['Nmesh'][i] for i in self.attrs['axes']], dtype='int')
        boxsize = numpy.array([self.attrs['BoxSize'][i] for i in self.attrs['axes']])
        I = numpy.eye(len(shape), dtype='int') * -2 + 1

        k = [numpy.fft.fftfreq(N, 1. / (N * 2 * numpy.pi / L))[:pkshape].reshape(kshape) for N, L, kshape, pkshape in zip(shape, boxsize, I, pk.shape)]

        kmag = sum(ki ** 2 for ki in k) ** 0.5
        W = numpy.empty(pk.shape, dtype='f4')
        W[...] = 2.0
        W[..., 0] = 1.0
        W[..., -1] = 1.0

        dk = self.attrs['dk']
        kmin = self.attrs['kmin']
        axes = list(self.attrs['axes'])
        kedges = numpy.arange(kmin, numpy.pi * self.attrs['Nmesh'][axes].min() / self.attrs['BoxSize'][axes].max() + dk/2, dk)

        xsum = numpy.zeros(len(kedges) + 1)
        Psum = numpy.zeros(len(kedges) + 1, dtype='complex128')
        Nsum = numpy.zeros(len(kedges) + 1)

        dig = numpy.digitize(kmag.flat, kedges)
        xsum.flat += numpy.bincount(dig, weights=(W * kmag).flat, minlength=xsum.size)
        Psum.real.flat += numpy.bincount(dig, weights=(W * pk.real).flat, minlength=xsum.size)
        Psum.imag.flat += numpy.bincount(dig, weights=(W * pk.imag).flat, minlength=xsum.size)
        Nsum.flat += numpy.bincount(dig, weights=W.flat, minlength=xsum.size)

        self.power = numpy.empty(len(kedges) - 1,
                dtype=[('k', 'f8'), ('power', 'c16'), ('modes', 'f8')])

        with numpy.errstate(invalid='ignore'):
            self.power['k'] = (xsum / Nsum)[1:-1]
            self.power['power'] = (Psum / Nsum)[1:-1] * boxsize.prod() # dimension is 'volume'
            self.power['modes'] = Nsum[1:-1]

        self.edges = kedges

        self.power = BinnedStatistic(['k'], [self.edges], self.power)

    def __getstate__(self):
        state = dict(
                     edges=self.edges,
                     power=self.power.data,
                     attrs=self.attrs)
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.power = BinnedStatistic(['k'], [self.edges], self.power)

def project_to_basis(y3d, edges, los=[0, 0, 1], poles=[]):
    """
    Project a 3D statistic on to the specified basis. The basis will be one
    of:

    - 2D (`x`, `mu`) bins: `mu` is the cosine of the angle to the line-of-sight
    - 2D (`x`, `ell`) bins: `ell` is the multipole number, which specifies
      the Legendre polynomial when weighting different `mu` bins

    .. note::

        The 2D (`x`, `mu`) bins will be computed only if `poles` is specified.
        See return types for further details.

    Notes
    -----
    *   the `mu` range extends from 0.0 to 1.0
    *   the `mu` bins are half-inclusive half-exclusive, except the last bin
        is inclusive on both ends (to include `mu = 1.0`)

    Parameters
    ----------
    y3d : RealField or ComplexField
        the 3D array holding the statistic to be projected to the specified basis
    edges : list of arrays, (2,)
        list of arrays specifying the edges of the desired `x` bins and `mu` bins
    los : array_like,
        the line-of-sight direction to use, which `mu` is defined with
        respect to; default is [0, 0, 1] for z.
    poles : list of int, optional
        if provided, a list of integers specifying multipole numbers to
        project the 2d `(x, mu)` bins on to
    hermitian_symmetric : bool, optional
        Whether the input array `y3d` is Hermitian-symmetric, i.e., the negative
        frequency terms are just the complex conjugates of the corresponding
        positive-frequency terms; if ``True``, the positive frequency terms
        will be explicitly double-counted to account for this symmetry

    Returns
    -------
    result : tuple
        the 2D binned results; a tuple of ``(xmean_2d, mumean_2d, y2d, N_2d)``, where:

        - xmean_2d : array_like, (Nx, Nmu)
            the mean `x` value in each 2D bin
        - mumean_2d : array_like, (Nx, Nmu)
            the mean `mu` value in each 2D bin
        - y2d : array_like, (Nx, Nmu)
            the mean `y3d` value in each 2D bin
        - N_2d : array_like, (Nx, Nmu)
            the number of values averaged in each 2D bin

    pole_result : tuple or `None`
        the multipole results; if `poles` supplied it is a tuple of ``(xmean_1d, poles, N_1d)``,
        where:

        - xmean_1d : array_like, (Nx,)
            the mean `x` value in each 1D multipole bin
        - poles : array_like, (Nell, Nx)
            the mean multipoles value in each 1D bin
        - N_1d : array_like, (Nx,)
            the number of values averaged in each 1D bin
    """
    comm = y3d.pm.comm
    x3d = y3d.x
    hermitian_symmetric = numpy.iscomplexobj(y3d)

    from scipy.special import legendre

    # setup the bin edges and number of bins
    xedges, muedges = edges
    x2edges = xedges**2
    Nx = len(xedges) - 1
    Nmu = len(muedges) - 1

    # always make sure first ell value is monopole, which
    # is just (x, mu) projection since legendre of ell=0 is 1
    do_poles = len(poles) > 0
    _poles = [0]+sorted(poles) if 0 not in poles else sorted(poles)
    legpoly = [legendre(l) for l in _poles]
    ell_idx = [_poles.index(l) for l in poles]
    Nell = len(_poles)

    # valid ell values
    if any(ell < 0 for ell in _poles):
        raise ValueError("in `project_to_basis`, multipole numbers must be non-negative integers")

    # initialize the binning arrays
    musum = numpy.zeros((Nx+2, Nmu+2))
    xsum = numpy.zeros((Nx+2, Nmu+2))
    ysum = numpy.zeros((Nell, Nx+2, Nmu+2), dtype=y3d.dtype) # extra dimension for multipoles
    Nsum = numpy.zeros((Nx+2, Nmu+2), dtype='i8')

    # if input array is Hermitian symmetric, only half of the last
    # axis is stored in `y3d`
    symmetry_axis = -1 if hermitian_symmetric else None

    # iterate over y-z planes of the coordinate mesh
    for slab in SlabIterator(x3d, axis=0, symmetry_axis=symmetry_axis):

        # the square of coordinate mesh norm
        # (either Fourier space k or configuraton space x)
        xslab = slab.norm2()

        # if empty, do nothing
        if len(xslab.flat) == 0: continue

        # get the bin indices for x on the slab
        dig_x = numpy.digitize(xslab.flat, x2edges)

        # make xslab just x
        xslab **= 0.5

        # get the bin indices for mu on the slab
        mu = slab.mu(los) # defined with respect to specified LOS
        dig_mu = numpy.digitize(abs(mu).flat, muedges)

        # make the multi-index
        multi_index = numpy.ravel_multi_index([dig_x, dig_mu], (Nx+2,Nmu+2))

        # sum up x in each bin (accounting for negative freqs)
        xslab[:] *= slab.hermitian_weights
        xsum.flat += numpy.bincount(multi_index, weights=xslab.flat, minlength=xsum.size)

        # count number of modes in each bin (accounting for negative freqs)
        Nslab = numpy.ones_like(xslab) * slab.hermitian_weights
        Nsum.flat += numpy.bincount(multi_index, weights=Nslab.flat, minlength=Nsum.size)

        # compute multipoles by weighting by Legendre(ell, mu)
        for iell, ell in enumerate(_poles):

            # weight the input 3D array by the appropriate Legendre polynomial
            weighted_y3d = legpoly[iell](mu) * y3d[slab.index]

            # add conjugate for this kx, ky, kz, corresponding to
            # the (-kx, -ky, -kz) --> need to make mu negative for conjugate
            # Below is identical to the sum of
            # Leg(ell)(+mu) * y3d[:, nonsingular]    (kx, ky, kz)
            # Leg(ell)(-mu) * y3d[:, nonsingular].conj()  (-kx, -ky, -kz)
            # or
            # weighted_y3d[:, nonsingular] += (-1)**ell * weighted_y3d[:, nonsingular].conj()
            # but numerically more accurate.
            if hermitian_symmetric:

                if ell % 2: # odd, real part cancels
                    weighted_y3d.real[slab.nonsingular] = 0.
                    weighted_y3d.imag[slab.nonsingular] *= 2.
                else:  # even, imag part cancels
                    weighted_y3d.real[slab.nonsingular] *= 2.
                    weighted_y3d.imag[slab.nonsingular] = 0.

            # sum up the weighted y in each bin
            weighted_y3d *= (2.*ell + 1.)
            ysum[iell,...].real.flat += numpy.bincount(multi_index, weights=weighted_y3d.real.flat, minlength=Nsum.size)
            if numpy.iscomplexobj(ysum):
                ysum[iell,...].imag.flat += numpy.bincount(multi_index, weights=weighted_y3d.imag.flat, minlength=Nsum.size)

        # sum up the absolute mag of mu in each bin (accounting for negative freqs)
        mu[:] *= slab.hermitian_weights
        musum.flat += numpy.bincount(multi_index, weights=abs(mu).flat, minlength=musum.size)

    # sum binning arrays across all ranks
    xsum  = comm.allreduce(xsum)
    musum = comm.allreduce(musum)
    ysum  = comm.allreduce(ysum)
    Nsum  = comm.allreduce(Nsum)

    # add the last 'internal' mu bin (mu == 1) to the last visible mu bin
    # this makes the last visible mu bin inclusive on both ends.
    ysum[..., -2] += ysum[..., -1]
    musum[:, -2]  += musum[:, -1]
    xsum[:, -2]   += xsum[:, -1]
    Nsum[:, -2]   += Nsum[:, -1]

    # reshape and slice to remove out of bounds points
    sl = slice(1, -1)
    with numpy.errstate(invalid='ignore'):

        # 2D binned results
        y2d       = (ysum[0,...] / Nsum)[sl,sl] # ell=0 is first index
        xmean_2d  = (xsum / Nsum)[sl,sl]
        mumean_2d = (musum / Nsum)[sl, sl]
        N_2d      = Nsum[sl,sl]

        # 1D multipole results (summing over mu (last) axis)
        if do_poles:
            N_1d     = Nsum[sl,sl].sum(axis=-1)
            xmean_1d = xsum[sl,sl].sum(axis=-1) / N_1d
            poles    = ysum[:, sl,sl].sum(axis=-1) / N_1d
            poles    = poles[ell_idx,...]

    # return y(x,mu) + (possibly empty) multipoles
    result = (xmean_2d, mumean_2d, y2d, N_2d)
    pole_result = (xmean_1d, poles, N_1d) if do_poles else None
    return result, pole_result

def _cast_source(source, BoxSize, Nmesh):
    """
    Cast an object to a MeshSource. BoxSize and Nmesh is used
    only on CatalogSource
    """
    from pmesh.pm import Field
    from nbodykit.source.mesh import FieldMesh

    if isinstance(source, Field):
        # if input is a Field object, wrap it as a MeshSource.
        source = FieldMesh(source)
    elif isinstance(source, CatalogSourceBase):
        # if input is CatalogSource, use defaults to make it into a mesh
        if not isinstance(source, MeshSource):
            source = source.to_mesh(BoxSize=BoxSize, Nmesh=Nmesh, dtype='f8', compensated=True)

    # now we are having a MeshSource
    # paint the density field to the mesh
    if not isinstance(source, MeshSource):
        raise TypeError("Unknown type of source in FFTPower: %s" % str(type(source)))
    if BoxSize is not None and any(source.attrs['BoxSize'] != BoxSize):
        raise ValueError("Mismatched Boxsize between __init__ and source.attrs")
    if Nmesh is not None and any(source.attrs['Nmesh'] != Nmesh):
        raise ValueError(("Mismatched Nmesh between __init__ and source.attrs; "
                          "if trying to re-sample with a different mesh, specify "
                          "`Nmesh` as keyword of to_mesh()"))

    return source
