# Source code for waveformtools.grids

    """ Classes to hold grid information

    Classes
    -------
    spherical_grid: grid info.
                    Stores information on the 2d grid in spherical polar coordinates.
    gl_grid : grid info
              Stores a Gauss-Legendre type grid on a spherical surface.
    """

    import numpy as np

    # from numba import jit, njit

    # import numba as nb

    # from numba.experimental import jitclass

    class spherical_grid:
        """A class to store the coordinate grid on a sphere.

        Attributes
        ----------
        ntheta: int
                                The number of angular points in the :math:`\\theta`
                                direction, including ghost zones.
        nphi:   int
                        The number of angular points in the :math:`\\phi`
                        direction, including ghost zones.
        nghosts:    int
                                The number of ghost zones at the end of
                                each direction.
        meshgrid:   tuple of 2d array
                                The 2d array containing the meshgrid of
                                (:math:`\\theta, \\phi`) angular points.
        theta_1d:   1d array
                                The 1d array of angular points
                                along the :math:`\\theta` axis.
        phi_1d: 1d array
                                The 1d array of angular points
                                along the :math:`\\phi` axis.
        dtheta: float
                                The angular step size in the :math:`\\theta`
                                direction.
        dphi:   float
                        The angular step size inthe :math:`\\phi`
                        direction.
        npix_act:   int
                                The total number of gridpoints on the sphere,
                                excluding the ghost points.
        meshgrid:
                                Get the 2d angular grid.

        Methods
        -------
        theta_1d:
                                Get the :math:`\\theta` axis.
        phi_1d:
                                Get the :math:`\\phi` axis.

        """

        def __init__(self, nphi=80, ntheta=41, nphimax=124, nthetamax=66, nghosts=2):
            # Number of gridpoints along phi direction including ghost points.
            self.nphi = nphi
            # Number of gridpoints along theta direction including ghost points.
            self.ntheta = ntheta
            # Total length of phi array used by ETK.
            self.nphimax = nphimax
            # Total length of theta array used by ETK.
            self.nthetamax = nthetamax
            # Number of ghost points in theta/phi direction.
            self.nghosts = nghosts

        @property
        def npix(self):
            """Return the total number of pixels,
            including the ghost zones present
            at one iteration."""
            return (self.ntheta) * (self.nphi)

        @property
        def npix_act(self):
            """Return the actual number of pixels,
            excluding the ghost zones present at
            one iteration"""
            return (self.ntheta - 2 * self.nghosts) * (self.nphi - 2 * self.nghosts)

        @property
        def npix_max(self):
            """Return the (max) total number of pixels,
            including the ghost and buffer zones at
            one iteration."""
            return (self.nthetamax) * (self.nphimax)

        @property
        def ntheta_act(self):
            """Return the actual number of valid pixels,
            excluding the ghost and buffer zones, along
            the theta axis at one iteration."""

            return self.ntheta - 2 * self.nghosts

        @property
        def nphi_act(self):
            """Return the actual number of valid pixels,
            excluding the ghost and buffer zones, along
            the phi axis at one iteration."""
            return self.nphi - 2 * self.nghosts

        @property
        def dtheta(self):
            """Return the coodinate spacing :math:`d\\theta.`"""
            return np.pi / (self.ntheta - 2 * self.nghosts)

        @property
        def dphi(self):
            """Return the coordinate spacing :math:`d\\phi`."""
            return 2 * np.pi / (self.nphi - 2 * self.nghosts)

        @property
        def nbuffer(self):
            """Return the number of buffer zones (excluding ghosts)"""
            return self.nthetamax - self.ntheta

        def theta_1d(self, theta_index=None):
            """Returns the coordinate value theta given
            the coordinate index. The coordinate index
            ranges from (0, ntheta). The actual indices
            without the ghost and extra zones
            is (nghosts, ntheta-nghosts).

            Parameters
            -----------

            theta_index:    int/ 1d array
                                            The theta coordinate index or axis.

            Returns
            -------

            theta_1d:   float
                                    The coordinate(s) :math:`\\theta` on the sphere.

            """

            if not theta_index:
                theta_index = np.arange(self.nghosts, self.ntheta - self.nghosts)

            return (theta_index - self.nghosts + 0.5) * np.pi / (self.ntheta - 2 * self.nghosts)

        def phi_1d(self, phi_index=None):
            """Returns the coordinate value theta
            given the coordinate index. The coordinate
            index lies in (0, nphi). The actual indices
            without the ghost and extra zones
            is (nghosts, nphi-nghosts).

            Parameters
            -----------

            phi_1d: int / 1d array
                                    The phi coordinate index or axis.

            Returns
            -------

            phi_1d: float or 1d array
                                    The coordinate(s) :math:`\\phi` on the sphere.

            """

            if not phi_index:
                phi_index = np.arange(self.nghosts, self.nphi - self.nghosts)

            return (phi_index - self.nghosts) * 2 * np.pi / (self.nphi - 2 * self.nghosts)

        @property
        def meshgrid(self):
            """The (:math:`\\theta, \\phi)`: coordinate meshes.
                    Excludes the ghost zones.

            Returns
            -------

            theta:  2d array
                            The :math:`\\theta` coordinate matrix for vectorization.

            phi:    2d array
                            The :math:`\\phi` coordinate matrix for vectorization.

            """

            theta, phi = np.meshgrid(self.theta_1d(), self.phi_1d())

            # theta = np.zeros((self.ntheta_act, self.nphi_act))
            # phi   = np.zeros((self.ntheta_act, self.nphi_act))

            # for theta_index, theta_val in enumerate(self.theta_1d()):
            #   for phi_index, phi_val in enumerate(self.phi_1d()):
            #       theta[theta_index, phi_index] = theta_val
            #       phi[theta_index, phi_index] = phi_val

            return np.transpose(theta), np.transpose(phi)

    class gl_grid:
        """A class to store the coordinate grid on a sphere.

        Attributes
        ----------
        ntheta: int
                    The number of angular points in the :math:`\\theta`
                    direction, including ghost zones.
        nphi:   int
                The number of angular points in the :math:`\\phi`
                direction, including ghost zones.
        nghosts:    int
                    The number of ghost zones at the end of
                    each direction.
        meshgrid:   tuple of 2d array
                    The 2d array containing the meshgrid of
                    (:math:`\\theta, \\phi`) angular points.
        theta_1d:   1d array
                    The 1d array of angular points
                    along the :math:`\\theta` axis.
        phi_1d: 1d array
                    The 1d array of angular points
                    along the :math:`\\phi` axis.
        dtheta: float
                    The angular step size in the :math:`\\theta`
                    direction.
        dphi:   float
                The angular step size inthe :math:`\\phi`
                direction.
        npix_act:   int
                    The total number of gridpoints on the sphere,
                    excluding the ghost points.
        meshgrid:
                    Get the 2d angular grid.

        Methods
        -------
        theta_1d:
                    Get the :math:`\\theta` axis.
        phi_1d:
                    Get the :math:`\\phi` axis.

        Notes
        -----

        The total number of points on the sphere is assumed to be :math:`(L+1)^2`

        :math:`N_\theta = L+1`

        :math:`N_\phi = 2(L+1)`

        This integrates out spherical harmonics of degree L exactly,
        given a regular function on the sphere.
        """

        def __init__(self, nphi=None, ntheta=None, L=47, nghosts=2):
            # Number of gridpoints along phi direction including ghost points.
            self._nphi = nphi
            # Number of gridpoints along theta direction including ghost points.
            self._ntheta = ntheta
            # Total length of phi array used by ETK.
            # self._nphimax   = nphimax
            # Total length of theta array used by ETK.
            # self._nthetamax = nthetamax
            # Number of ghost points in theta/phi direction.
            self._nghosts = nghosts
            self._L = L
            self._theta_1d = None
            self._phi_1d = None
            self._meshgrid = None

            if self._ntheta is None:
                if self._L is None:
                    raise ValueError("Please specify L or angular points!")

                else:
                    self._ntheta_act = L + 1
                    self._nphi_act = 2 * self._ntheta_act

                    self._ntheta = self._ntheta_act + 2 * self._nghosts
                    self._nphi = self._nphi_act + 2 * self._nghosts

            elif self._L is None:
                if self.nthetha is None:
                    raise ValueError("Please specify L or angular points!")
                else:
                    self.L = self.ntheta - 1
                    assert nphi % 2 == 0

            from scipy.special import roots_legendre

            cpoints, self._weights, self._sum_of_weights = roots_legendre(L + 1, mu=True)

            # xpoints = (np.pi-np.arccos(cpoints))
            xpoints = np.arccos(cpoints[::-1])
            self._theta_1d = xpoints

            dphi = 2 * np.pi / self._nphi_act

            self._phi_1d = np.linspace(0, 2 * np.pi - dphi, self._nphi_act)

            theta_grid, phi_grid = np.meshgrid(self._theta_1d, self._phi_1d)
            self._meshgrid = np.transpose(theta_grid), np.transpose(phi_grid)

            dtheta_axis = np.diff(self._theta_1d)

            dtheta_axis = np.append(dtheta_axis, np.pi - self._theta_1d[-1])
            dtheta_axis = np.insert(dtheta_axis, 0, self._theta_1d[0])

            self._dtheta_1d = dtheta_axis

            self._dphi = dphi  # self._phi_1d[1]

        @property
        def nphi(self):
            """Return the (max) total number of pixels,
            including the ghost and buffer zones at
            one iteration."""
            return self._nphi

        @property
        def ntheta(self):
            """Return the actual number of valid pixels,
            excluding the ghost and buffer zones, along
            the theta axis at one iteration."""
            return self._ntheta

        @property
        def nghosts(self):
            """Return the number of ghost zones at each end in each direction"""
            return self._nghosts

        @property
        def ntheta_act(self):
            """Return the actual number of valid pixels,
            excluding the ghost and buffer zones, along
            the phi axis at one iteration."""
            return self.ntheta - 2 * self.nghosts

        @property
        def nphi_act(self):
            """Return the actual number of valid pixels,
            excluding the ghost and buffer zones, along
            the phi axis at one iteration."""
            return self.nphi - 2 * self.nghosts

        @property
        def npix_act(self):
            """Return the actual number of pixels,
            excluding the ghost zones present at
            one iteration"""
            return (self.ntheta_act) * (self.nphi_act)

        @property
        def dtheta_1d(self):
            """Return the non-uniform angular stepping
            in :math:`\theta` direction"""
            return self._dtheta_1d

        @property
        def dphi(self):
            """Return the uniform angular stepping
            in :math:`\phi` direction"""
            return self._dphi

        @property
        def L(self):
            """Return the maximum mode number
            that this grid can integrate out."""

            return self._L

        @property
        def npix(self):
            """Return the total number of pixels,
            including the ghost zones present at
            one iteration."""
            return (self.ntheta) * (self.nphi)

        @property
        def weights(self):
            """Return the integration weights along theta."""
            return self._weights

        @property
        def weights_grid(self):
            """Return the integration weights on the coordinate mesh grid"""
            return np.outer(self.weights, np.ones(self.nphi_act))

        @property
        def shape(self):
            """Return the shape of the grid."""
            return self.weights_grid.shape

        @property
        def theta_1d(self):
            """Returns the coordinate values theta given the coordinate index.
            The coordinate index ranges from (0, ntheta). The actual indices
            without the ghost and extra zones is (nghosts, ntheta-nghosts).

            Parameters
            -----------

            theta_index:    int/ 1d array
                            The theta coordinate index or axis.

            Returns
            -------

            theta_1d:   float
                        The coordinate(s) :math:`\\theta` on the sphere.

            """

            return self._theta_1d

        @property
        def phi_1d(self):
            """Returns the coordinate values phi given the coordinate index.
            The coordinate index lies in (0, nphi). The actual indices without
            the ghost and extra zones is (nghosts, nphi-nghosts).

            Parameters
            -----------

            phi_1d: int / 1d array
                        The phi coordinate index or axis.

            Returns
            -------

            phi_1d: float or 1d array
                        The coordinate(s) :math:`\\phi` on the sphere.

            """

            return self._phi_1d

        @property
        def meshgrid(self):
            """The (:math:`\\theta, \\phi)`: coordinate meshes.
                Excludes the ghost zones.

            Returns
            -------

            theta:  2d array
                    The :math:`\\theta` coordinate matrix for vectorization.

            phi:    2d array
                    The :math:`\\phi` coordinate matrix for vectorization.

            """

            return self._meshgrid