# Source code for waveformtools.spherical

    """ Methods to handle functions on a sphere. """

    ##################################
    # Imports
    #################################

    import numpy as np

    def decompose_in_SWSHs(waveform, gridinfo, spin_weight=-2, ell_max=8, emm_list="all"):
        """Decompose a given function on a sphere in Spin Weighted Spherical Harmonics

        Parameters
        ----------

        waveform:		list
                        A list that contains as its items the waveform
                        defined on the sphere as an array of shape
                        [ntheta, nphi]. Each item in the list may denote
                        an instant of time or frequency.

        spin_weight:	 int, optional
                        The spin weight of the waveform.
                        It defaults to -2 for a gravitational waveform.

        ell_max:	int, optional
                    The maximum value of the :math:`\\ell'
                    polar quantum number. Defaults to 6=8.

        gridinfo:		class instance
                        The class instance that contains
                        the properties of the spherical grid.

        Returns
        -------

        SWSH_coeffs:	list
                        The SWSH coefficients of the waveform.
                        It may be a list composed of a single
                        floating point number or a 1d array
                        (denoting time or frequency dimension).
                        The waveform can have angular as well
                        as time dimentions. The nesting order
                        will be that, given the list
                        `non_boosted_waveform', each
                        item refers to a one dimensional array in
                        time/ frequency of SWSH coefficients.

        Notes
        -----

        Assumes that the sphere on which this decomposition is carried out is so far out
        that the coordinate system is spherical polar and the poper area is the
        same as its co-ordinate area.

        """

        # Find out if the unboosted waveform is a single number or defined on a spherical grid.
        onepoint = isinstance(waveform[0], float)

        if not onepoint:
            # Get the spherical grid shape.
            ntheta, nphi = np.array(waveform[0]).shape

        # Compute the meshgrid for theta and phi.
        theta = gridinfo.theta(ntheta=ntheta, nphi=nphi)

        # Compute the meshgrid for theta and phi.
        phi = gridinfo.phi(ntheta=ntheta, nphi=nphi)

        # decomposed_waveforms = {}

        multipoles_all = {}
        for item in waveform:
            # Integrate on the sphere for decomposition into SWSHs

            """This should be fixed with summation over ell"""

            # define m values.
            if emm_list == "all":
                emm_list = np.arange(-ell_max, ell_max + 1)

            # Convert input to arrays.

            integrand_data = np.array(item)

            # Step 1: Compute the surface integral.

            # Assign data vectors.
            # lpole  = {}

            # Check if data includes ghost zones or not.

            # Compute the meshgrid for theta and phi.
            theta = gridinfo.theta(ntheta=ntheta, nphi=nphi)

            # Compute the meshgrid for theta and phi.
            phi = gridinfo.phi(ntheta=ntheta, nphi=nphi)

            sqrt_met_det = np.sqrt(np.power(np.sin(theta), 2))

            integrand_ij = integrand_data

            darea = sqrt_met_det * gridinfo.dtheta * gridinfo.dphi

            from waveformtools.transforms import Yslm_vec

            for ell_index in range(ell_max):
                multipoles_ell = {}
                for emm_index in range(len(emm_list)):
                    # Decompose into seperate m modes.

                    # m value.
                    emm_val = int(emm_list[emm_index])

                    # Spin weighted spherical harmonic function at (theta, phi)
                    Ybasis_fun = Yslm_vec(spin_weight, ell_index, emm_val, theta, phi)

                    # Integrate to obtain the multipole of order l.

                    # Integration for real and imaginary parts of the data separately.
                    # Integrate the function

                    # Using quad
                    multipole_emm = quad_on_sphere(integrand_ij * Ybasis_fun * darea, gridinfo)
                    # multipole_emm	 =	 np.sum(integrand_ij * Ybasis_fun * darea)

                    multipoles_ell.update({emm_val: multipole_emm})
                multipoles_all.update({ell_index: multipoles_ell})

            # Return the computed multipole.
            return multipoles_all

    def quad_on_sphere(integrand, gridinfo, kind="third"):
        """Integrate on a sphere using the scipy.quad method

        Parameters
        ----------

        integrand:		2d array
                                        The two dimensional integrand array defined on the sphere.

        info:		class instance
                                The class instance that contains the properties of the spherical grid.

        kind:		str
                                The interpolation order to use in integration.
        Returns
        -------

        final_integral : float
                                  The given integrand integrated over the sphere.

        final_errs:	float
                                        The accumulated errors.

        Notes
        -----

        Assumes that the sphere is a unit round sphere.

        """

        # Step 0: Get the grid properties

        # Compute the meshgrid for theta and phi.
        theta_1d = gridinfo.theta_1d
        # theta
        # Compute the meshgrid for theta and phi.
        phi_1d = gridinfo.phi_1d

        # imports
        from scipy.interpolate import interp1d
        from scipy.integrate import quad

        theta_first_integral_vals = []
        theta_first_integral_errs = []
        # Step 1: integrate along the theta direction
        for phi_index in range(gridinfo.nphi):
            # Interpolate the integrand.

            integrand_phi = integrand[:, phi_index]

            integrand_phi_interp_func = interp1d(theta_1d, integrand_phi, kind=kind)

            # Integrate on the phi plane
            integral_phi_vals, integral_phi_errs = quad(integrand_phi_interp_func, 0, np.pi)

            theta_first_integral_vals.append(integral_phi_vals)
            theta_first_integral_errs.append(integral_phi_errs)

        # Step 2: integrate along the phi direction

        # Interpolate the integrand.

        integrand_theta = theta_first_integral_vals

        integrand_theta_interp_func = interp1d(phi_1d, integrand_theta, kind=kind)

        # Integrate on the theta plane
        final_integral, semi_final_errs = quad(integrand_theta_interp_func, 0, 2 * np.pi)

        # Get final errors
        final_errs = semi_final_errs + np.sum(np.array(theta_first_integral_errs))

        return final_integral, final_errs