# coding: utf-8
# Wrapper around the query_disc method of Healpix_base class

import numpy as np
cimport numpy as np
from libcpp.string cimport string
from libcpp.vector cimport vector
from libc.math cimport sqrt, floor, fabs
cimport libc
from . import npix2nside, nside2npix
from .pixelfunc import maptype, pix2ang
import os
import cython
from libcpp cimport bool as cbool

from _common cimport tsize, arr, xcomplex, Healpix_Ordering_Scheme, RING, NEST, Healpix_Map, Alm, ndarray2map, ndarray2alm, rotmatrix

from . import _healpy_sph_transform_lib as sphtlib

cdef double UNSEEN = -1.6375e30
cdef double rtol_UNSEEN = 1.e-7 * 1.6375e30

cdef extern from "alm_healpix_tools.h":
    cdef void map2alm_iter(Healpix_Map[double] &m,
                           Alm[xcomplex[double]] &alm,
                           int num_iter, arr[double] &weight)
    cdef void map2alm_pol_iter(Healpix_Map[double] &mapT,
                               Healpix_Map[double] &mapQ,
                               Healpix_Map[double] &mapU,
                               Alm[xcomplex[double]] &almT,
                               Alm[xcomplex[double]] &almG,
                               Alm[xcomplex[double]] &almC,
                               int num_iter,
                               arr[double] &weight)
    cdef void map2alm_spin(    Healpix_Map[double] &map1,
                               Healpix_Map[double] &map2,
                               Alm[xcomplex[double]] &alm1,
                               Alm[xcomplex[double]] &alm2,
                               int spin,
                               arr[double] &weight,
                               cbool add_alm)
    cdef void alm2map_spin(    Alm[xcomplex[double]] &alm1,
                               Alm[xcomplex[double]] &alm2,
                               Healpix_Map[double] &map1,
                               Healpix_Map[double] &map2,
                               int spin)

cdef extern from "alm_powspec_tools.h":
    cdef void c_rotate_alm "rotate_alm" (Alm[xcomplex[double]] &alm, double psi, double theta, double phi)
    cdef void c_rotate_alm "rotate_alm" (Alm[xcomplex[double]] &ai, Alm[xcomplex[double]] &ag, Alm[xcomplex[double]] &ac, double psi, double theta, double phi)
    cdef void c_rotate_alm "rotate_alm" (Alm[xcomplex[double]] &alm, rotmatrix &mat)
    cdef void c_rotate_alm "rotate_alm" (Alm[xcomplex[double]] &ai, Alm[xcomplex[double]] &ag, Alm[xcomplex[double]] &ac, rotmatrix &mat)

cdef extern from "healpix_data_io.h":
    cdef void read_weight_ring (string &dir, int nside, arr[double] &weight)
    cdef vector[double] read_fullweights_from_fits(string &weightfile, int nside)

cdef extern from "weight_utils.h":
    cdef void apply_fullweights (Healpix_Map[double] &map,
            vector[double] &wgt)

DATAPATH = None

def get_datapath():
    global DATAPATH
    if DATAPATH is None:
        DATAPATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
    return DATAPATH

def map2alm_spin_healpy(maps, spin, lmax = None, mmax = None):
    """Computes the spinned alm of a 2 Healpix maps.

    Parameters
    ----------
    m : list of 2 arrays
        list of 2 input maps as numpy arrays
    spin : int
        spin of the alms
    lmax : int, scalar, optional
      Maximum l of the power spectrum. Default: 3*nside-1
    mmax : int, scalar, optional
      Maximum m of the alm. Default: lmax

    Returns
    -------
    alms : list of 2 arrays
      list of 2 alms

    Warnings
    --------
    Previously, ``spin=0`` was an invalid parameter value, and silently
    returned undefined results.

    """
    spin = int(spin)

    if spin == 0:
        return [-map2alm(mm, niter=0, lmax=lmax, mmax=mmax) for mm in maps]

    maps_c = [np.ascontiguousarray(m, dtype=np.float64) for m in maps]

    # create UNSEEN mask for map
    masks = [mkmask(m) for m in maps_c]

    # Adjust lmax and mmax
    cdef int lmax_, mmax_, nside, npix
    npix = maps_c[0].size
    nside = npix2nside(npix)
    if lmax is None:
        lmax_ = 3 * nside - 1
    else:
        lmax_ = lmax
    if mmax is None:
        mmax_ = lmax_
    else:
        mmax_ = mmax

    # Check all maps have same npix
    if maps_c[1].size != npix:
        raise ValueError("Input maps must have same size")

    # View the ndarray as a Healpix_Map
    M1 = ndarray2map(maps_c[0], RING)
    M2 = ndarray2map(maps_c[1], RING)

    # replace UNSEEN pixels with zeros
    for m, mask in zip(maps_c, masks):
        if mask is not False:
            m[mask] = 0.0

    # Create an ndarray object that will contain the alm for output (to be returned)
    n_alm = alm_getn(lmax_, mmax_)
    alms = [np.empty(n_alm, dtype=np.complex128) for m in maps]

    # View the ndarray as an Alm
    # Alms = [ndarray2alm(alm, lmax_, mmax_) for alm in alms]
    A1 = ndarray2alm(alms[0], lmax_, mmax_)
    A2 = ndarray2alm(alms[1], lmax_, mmax_)

    # ring weights
    cdef arr[double] * w_arr = new arr[double]()
    cdef int i
    cdef char *c_datapath
    w_arr.allocAndFill(2 * nside, 1.)

    map2alm_spin(M1[0], M2[0], A1[0], A2[0], spin, w_arr[0], False)

    # restore input map with UNSEEN pixels
    for m, mask in zip(maps_c, masks):
        if mask is not False:
            m[mask] = UNSEEN

    del w_arr
    del M1, M2, A1, A2
    return alms

def alm2map_spin_healpy(alms, nside, spin, lmax, mmax=None):
    """Computes maps from a set of 2 spinned alm

    Parameters
    ----------
    alms : list of 2 arrays
      list of 2 alms
    nside : int
        requested nside of the output map
    spin : int
        spin of the alms
    lmax : int, scalar
      Maximum l of the power spectrum.
    mmax : int, scalar, optional
      Maximum m of the alm. Default: lmax

    Returns
    -------
    m : list of 2 arrays
        list of 2 out maps in RING scheme as numpy arrays

    Warnings
    --------
    Previously, ``spin=0`` was an invalid parameter value, and silently
    returned undefined results.

    """
    spin = int(spin)

    if spin == 0:
        if not mmax:
            mmax = lmax
        return [-sphtlib._alm2map(alm, nside, lmax=lmax, mmax=mmax) for alm in alms]

    alms_c = [np.ascontiguousarray(alm, dtype=np.complex128) for alm in alms]

    npix = nside2npix(nside)
    maps = [np.zeros(npix, dtype=np.float64) for alm in alms]

    # View the ndarray as a Healpix_Map
    M1 = ndarray2map(maps[0], RING)
    M2 = ndarray2map(maps[1], RING)

    if not mmax:
        mmax = lmax

    # View the ndarray as an Alm
    A1 = ndarray2alm(alms_c[0], lmax, mmax)
    A2 = ndarray2alm(alms_c[1], lmax, mmax)

    alm2map_spin(A1[0], A2[0], M1[0], M2[0], spin)

    del M1, M2, A1, A2
    return maps

def map2alm(m, lmax = None, mmax = None, niter = 3, use_weights = False,
            datapath = None, gal_cut = 0, pixel_weights_filename=None):
    """Computes the alm of a Healpix map.

    Parameters
    ----------
    m : array-like, shape (Npix,) or (3, Npix)
      The input map or a list of 3 input maps (polarization).
    lmax : int, scalar, optional
      Maximum l of the power spectrum. Default: 3*nside-1
    mmax : int, scalar, optional
      Maximum m of the alm. Default: lmax
    iter : int, scalar, optional
      Number of iteration (default: 1)
    use_weights: bool, scalar, optional
      If True, use the ring weighting. Default: False.
    gal_cut : float [degrees]
      pixels at latitude in [-gal_cut;+gal_cut] are not taken into account
    pixel_weights_filename : str
      filename of the pixel by pixel weights

    Returns
    -------
    alm : array or tuple of arrays
      alm or a tuple of 3 alm (almT, almE, almB) if polarized input.
    """
    # Check if the input map is polarized or not
    info = maptype(m)
    if info == 0:
        polarization = False
        mi = m.astype(np.float64, order='C', copy=True)
    elif info == 1:
        polarization = False
        mi = m[0].astype(np.float64, order='C', copy=True)
    elif info == 3:
        polarization = True
        mi = m[0].astype(np.float64, order='C', copy=True)
        mq = m[1].astype(np.float64, order='C', copy=True)
        mu = m[2].astype(np.float64, order='C', copy=True)
    else:
        raise ValueError("Wrong input map (must be a valid healpix map "
                         "or a sequence of 1 or 3 maps)")

    # replace UNSEEN pixels with zeros
    mask = mkmask(mi)
    if mask is not False:
        mi[mask] = 0
    if polarization:
        mask = mkmask(mq)
        if mask is not False:
            mq[mask] = 0
        mask = mkmask(mu)
        if mask is not False:
            mu[mask] = 0

    cdef int nside, npix
    npix = mi.size
    nside = npix2nside(npix)

    # Optionally apply a galactic cut
    if gal_cut is not None and gal_cut > 0:
        mask_gal = pix2ang(nside, np.arange(npix), lonlat=True)[1]
        mask_gal = np.abs(mask_gal) < gal_cut
        mi[mask_gal] = 0
        if polarization:
            mq[mask_gal] = 0
            mu[mask_gal] = 0
        del mask_gal

    # Adjust lmax and mmax
    cdef int lmax_, mmax_
    if lmax is None:
        lmax_ = 3 * nside - 1
    else:
        lmax_ = lmax
    if mmax is None:
        mmax_ = lmax_
    else:
        mmax_ = mmax

    # Check all maps have same npix
    if polarization:
        if mq.size != npix or mu.size != npix:
            raise ValueError("Input maps must have same size")

    # View the ndarray as a Healpix_Map
    MI = ndarray2map(mi, RING)
    if polarization:
        MQ = ndarray2map(mq, RING)
        MU = ndarray2map(mu, RING)


    # Create an ndarray object that will contain the alm for output (to be returned)
    n_alm = alm_getn(lmax_, mmax_)
    almI = np.empty(n_alm, dtype=np.complex128)
    if polarization:
        almG = np.empty(n_alm, dtype=np.complex128)
        almC = np.empty(n_alm, dtype=np.complex128)

    # View the ndarray as an Alm
    AI = ndarray2alm(almI, lmax_, mmax_)
    if polarization:
        AG = ndarray2alm(almG, lmax_, mmax_)
        AC = ndarray2alm(almC, lmax_, mmax_)

    # ring weights
    cdef arr[double] * w_arr = new arr[double]()
    cdef int i
    cdef char *c_datapath
    if use_weights:
        if datapath is None:
            datapath = get_datapath()
        # For Python3: datapath must be a str, bdatapath must be bytes
        if isinstance(datapath, unicode) :
            bdatapath = datapath.encode('UTF-8')
        else :
            bdatapath = datapath
            datapath = datapath.decode('UTF-8')
        c_datapath = bdatapath
        weightfile = 'weight_ring_n%05d.fits' % (nside)
        if not os.path.isfile(os.path.join(datapath, weightfile)):
            raise IOError('Weight file not found in %s' % (datapath))
        read_weight_ring(string(c_datapath), nside, w_arr[0])
        for i in range(w_arr.size()):
            w_arr[0][i] += 1
    else:
        w_arr.allocAndFill(2 * nside, 1.)

    cdef vector[double] full_weights
    if pixel_weights_filename is not None:

        full_weights = read_fullweights_from_fits(pixel_weights_filename.encode("UTF-8"), nside)
        # pixel weighting requires 0 iterations
        niter = 0
        apply_fullweights(MI[0], full_weights)
        if polarization:
            apply_fullweights(MQ[0], full_weights)
            apply_fullweights(MU[0], full_weights)

    if polarization:
        map2alm_pol_iter(MI[0], MQ[0], MU[0], AI[0], AG[0], AC[0], niter, w_arr[0])
    else:
        map2alm_iter(MI[0], AI[0], niter, w_arr[0])

    del w_arr
    if polarization:
        del MI, MQ, MU, AI, AG, AC
        return np.array([almI, almG, almC])
    else:
        del MI, AI
        return almI


def alm2cl(alms, alms2 = None, lmax = None, mmax = None, lmax_out = None):
    """Computes (cross-)spectra from alm(s). If alm2 is given, cross-spectra between
    alm and alm2 are computed. If alm (and alm2 if provided) contains n alm,
    then n(n+1)/2 auto and cross-spectra are returned.

    Parameters
    ----------
    alms : complex, array or sequence of arrays
      The alm from which to compute the power spectrum. If n>=2 arrays are given,
      computes both auto- and cross-spectra.
    alms2 : complex, array or sequence of 3 arrays, optional
      If provided, computes cross-spectra between alm and alm2.
      Default: alm2=alm, so auto-spectra are computed.
    lmax : None or int, optional
      The maximum l of the input alm. Default: computed from size of alm
      and mmax_in
    mmax : None or int, optional
      The maximum m of the input alm. Default: assume mmax_in = lmax_in
    lmax_out : None or int, optional
      The maximum l of the returned spectra. By default: the lmax of the given
      alm(s).

    Returns
    -------
    cl : array or tuple of n(n+1)/2 arrays
      the spectrum <*alm* x *alm2*> if *alm* (and *alm2*) is one alm, or
      the auto- and cross-spectra <*alm*[i] x *alm2*[j]> if alm (and alm2)
      contains more than one spectra.
      If more than one spectrum is returned, they are ordered by diagonal.
      For example, if *alm* is almT, almE, almB, then the returned spectra are:
      TT, EE, BB, TE, EB, TB.
    """
    #############################
    # Check alm and number of spectra
    #
    cdef int Nspec, Nspec2
    if not hasattr(alms, '__len__'):
        raise ValueError('alms must be an array or a sequence of arrays')
    if not hasattr(alms[0], '__len__'):
        alms_lonely = True
        alms = [alms]
    else:
        alms_lonely = False

    Nspec = len(alms)

    if alms2 is None:
        alms2 = alms

    if not hasattr(alms2, '__len__'):
        raise ValueError('alms2 must be an array or a sequence of arrays')
    if not hasattr(alms2[0], '__len__'):
        alms2 = [alms2]
    Nspec2 = len(alms2)

    if Nspec != Nspec2:
        raise ValueError('alms and alms2 must have same number of spectra')

    ##############################################
    # Check sizes of alm's and lmax/mmax/lmax_out
    #
    cdef int almsize
    almsize = alms[0].size
    for i in xrange(Nspec):
        if alms[i].size != almsize or alms2[i].size != almsize:
            raise ValueError('all alms must have same size')

    lmax, mmax = alm_getlmmax(alms[0], lmax, mmax)

    if lmax_out is None:
        lmax_out = lmax


    #######################
    # Computing the spectra
    #
    cdef int j, l, m, limit
    cdef int lmax_ = lmax, mmax_ = mmax
    cdef int lmax_out_ = lmax_out

    cdef np.ndarray[double, ndim=1] powspec_
    cdef np.ndarray[np.complex128_t, ndim=1] alm1_
    cdef np.ndarray[np.complex128_t, ndim=1] alm2_

    spectra = []
    for n in xrange(Nspec): # diagonal rank
        for m in xrange(0, Nspec - n): # position in the diagonal
            powspec_ = np.zeros(lmax + 1)
            alm1_ = alms[m]
            alm2_ = alms2[m + n]
            # compute cross-spectrum alm1[n] x alm2[n+m]
            # and place result in result list
            for l in range(lmax_ + 1):
                j = alm_getidx(lmax_, l, 0)
                powspec_[l] = alm1_[j].real * alm2_[j].real
                limit = l if l <= mmax else mmax
                for m in range(1, limit + 1):
                    j = alm_getidx(lmax_, l, m)
                    powspec_[l] += 2 * (alm1_[j].real * alm2_[j].real +
                                        alm1_[j].imag * alm2_[j].imag)
                powspec_[l] /= (2 * l + 1)
            spectra.append(powspec_[:lmax_out+1])

    # if only one alm was given, returns only cl and not a list with one cl
    if alms_lonely:
        spectra = spectra[0]

    return np.array(spectra)


@cython.wraparound(False)
@cython.boundscheck(False)
def almxfl(alm, fl, mmax = None, inplace = False):
    """Multiply an a_lm by a vector b_l.

    Parameters
    ----------
    alm : array, double
      The array representing the spherical harmonics coefficients
    fl : array, double
      The array giving the factor f_l by which to multiply a_lm
    mmax : None or int, optional
      The maximum m of the input alm
    inplace : bool, optional
      If True, performs the computation in-place if possible (input alm
      is modified if it is a 1d-array of type float64). Otherwise,
      a copy of alm is done.

    Returns
    -------
    alm : array, double
      The result of a_lm * f_l. If *inplace* is True, returns the input
      alm modified
    """
    cdef np.ndarray[np.complex128_t, ndim=1] alm_
    cdef np.ndarray[np.complex128_t, ndim=1] fl_

    if inplace:
        alm_ = np.ascontiguousarray(alm, dtype = np.complex128)
    else:
        alm_ = np.array(alm, dtype = np.complex128, copy = True)

    fl_ = np.ascontiguousarray(fl, dtype = np.complex128)

    cdef int lmax_, mmax_
    cdef int l, m
    lmax_, mmax_ = alm_getlmmax(alm_, None, mmax)

    cdef np.complex128_t f
    cdef int maxm, i
    cdef int flsize = fl_.size
    for l in xrange(lmax_ + 1):
        f = fl_[l] if l < flsize else 0.
        maxm = l if l <= mmax_ else mmax_
        for m in xrange(maxm + 1):
            i = alm_getidx(lmax_, l, m)
            alm_[i] *= f

    return alm_


def rotate_alm(alm not None, double psi=0, double theta=0, double phi=0, matrix=None, lmax=None,
               mmax=None):
    """
    This routine transforms the scalar (and tensor) a_lm coefficients
    to emulate the effect of an arbitrary rotation of the underlying
    map. The rotation is done directly on the a_lm using the Wigner
    rotation matrices, computed by recursion. To rotate the a_lm for
    l ≤ l_max the number of operations scales like l_max^3.

    Parameters
    ----------
    alm : array-like of shape (n,) or (k,n), or list of arrays
        Complex a_lm values before and after rotation of the coordinate system.
    psi : float
        First rotation: angle ψ about the z-axis. All angles are in radians
        and should lie in [-2pi,2pi], the rotations are active and the
        referential system is assumed to be right handed. The routine
        coordsys2euler zyz can be used to generate the Euler angles ψ, θ, φ
        for rotation between standard astronomical coordinate systems.
    theta : float
        Second rotation: angle θ about the original (unrotated) y-axis
    phi : float.
        Third rotation: angle φ about the original (unrotated) z-axis.
    matrix : np.ndarray
        Rotation matrix instead of angles, for example the output of hp.Rotator.mat
    lmax : int
        Maximum multipole order l of the data set.
    mmax : int
        Maximum degree m of data set.

    """
    if isinstance(alm, np.ndarray) and alm.ndim == 1:
        alm = [alm]

    cdef rotmatrix rotation_matrix
    if matrix is not None:
        assert matrix.shape == (3,3), "Rotation matrix should be a 3x3 array"

        rotation_matrix = rotmatrix(matrix[0,0], matrix[0,1], matrix[0,2],
                                    matrix[1,0], matrix[1,1], matrix[1,2],
                                    matrix[2,0], matrix[2,1], matrix[2,2])

    if not isinstance(alm, (list, tuple, np.ndarray)) or len(alm) == 0:
        raise ValueError('Invalid input.')

    if alm[0].dtype != np.complex128:
        raise ValueError('rotate_alm only supports double precision complex arrays')

    # C++ rotate_alm only handles 1 or 3 maps. The function handling 3 maps
    # is faster than running 3 times the 1-map function, but gives identical
    # results.
    if len(alm) not in (1, 3):
        for a in alm:
            if matrix is None:
                rotate_alm(a, psi, theta, phi)
            else:
                rotate_alm(a, matrix)
        return

    lmax, mmax = alm_getlmmax(alm[0], lmax, mmax)
    ai = np.ascontiguousarray(alm[0], dtype=np.complex128)
    AI = ndarray2alm(ai, lmax, mmax)
    if len(alm) == 1:
        if matrix is None:
            c_rotate_alm(AI[0], psi, theta, phi)
        else:
            c_rotate_alm(AI[0], rotation_matrix)
        del AI
    else:
        ag = np.ascontiguousarray(alm[1], dtype=np.complex128)
        ac = np.ascontiguousarray(alm[2], dtype=np.complex128)
        AG = ndarray2alm(ag, lmax, mmax)
        AC = ndarray2alm(ac, lmax, mmax)
        if matrix is None:
            c_rotate_alm(AI[0], AG[0], AC[0], psi, theta, phi)
        else:
            c_rotate_alm(AI[0], AG[0], AC[0], rotation_matrix)
        del AI, AG, AC


@cython.cdivision(True)
cdef int alm_getn(int l, int m):
    if not m <= l:
        raise ValueError("mmax must be <= lmax")
    return ((m+1)*(m+2))/2 + (m+1)*(l-m)


def alm_getlmmax(a, lmax, mmax):
    if lmax is None:
        if mmax is None:
            lmax = alm_getlmax(a.size)
            mmax = lmax
        else:
            lmax = alm_getlmax2(a.size, mmax)
    elif mmax is None:
        mmax = lmax
    return lmax, mmax


@cython.cdivision(True)
cdef inline int alm_getlmax(int s):
    cdef double x
    x=(-3+np.sqrt(1.+8.*s))/2
    if x != floor(x):
        return -1
    else:
        return <int>floor(x)


@cython.cdivision(True)
cdef inline int alm_getlmax2(int s, int mmax):
    cdef double x
    x = (2 * s + mmax ** 2 - mmax - 2.) / (2 * mmax + 2.)
    if x != floor(x):
        return -1
    else:
        return <int>floor(x)


@cython.cdivision(True)
cdef inline int alm_getidx(int lmax, int l, int m):
    return m*(2*lmax+1-m)/2+l


@cython.wraparound(False)
@cython.boundscheck(False)
cpdef mkmask(np.ndarray[double, ndim=1] m):
    cdef int nbad
    cdef int size = m.size
    cdef int i
    # first, count number of bad pixels, to see if allocating a mask is needed
    nbad = count_bad(m)
    cdef np.ndarray[np.int8_t, ndim=1] mask
    #cdef np.ndarray[double, ndim=1] m_
    if nbad == 0:
        return False
    else:
        mask = np.zeros(size, dtype = np.int8)
        #m_ = m
        for i in range(size):
            if fabs(m[i] - UNSEEN) < rtol_UNSEEN:
                mask[i] = 1
    mask.dtype = bool
    return mask

@cython.wraparound(False)
@cython.boundscheck(False)
cpdef int count_bad(np.ndarray[double, ndim=1] m):
    cdef int i
    cdef int nbad = 0
    cdef size = m.size
    for i in xrange(m.size):
        if fabs(m[i] - UNSEEN) < rtol_UNSEEN:
            nbad += 1
    return nbad
