# Source code for waveformtools.BMS

    """ The implementation of BMS transformations on the waveforms. """

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

    import numpy as np

    def compute_conformal_k(vec_v, theta, phi, spin_phase=0):
        """Compute the conformal factor for the boost transformation
                :math:`k = \\exp(-2i \\lambda) \\gamma^3
                (1 - \\mathbf{v} \\cdot \\mathbf{r})^3`

        Parameters
        ----------
        vec_v: list
                        The velocity vector.

        theta: float
                        The polar angle :math:`\\theta' in radians.

        phi: float
                                The azimuthal angle :math:`\\phi' in radians.

        spin_phase: float, optional
                                 The spin phase :math:`\\lambda'. Defaults to 0.

        Returns
        -------
        conformal_k:	float
                        The conformal factor for the
                        boost transformation as defined above.
        """

        # unpack the velocity vector
        vel_x, vel_y, vel_z = vec_v

        # magnitude of velocity
        mag_v = np.sqrt(vel_x**2 + vel_y**2 + vel_z**2)
        # compute the dot product
        v_dot_r = np.sin(theta) * (vel_x * np.cos(phi) + vel_y * np.sin(phi)) + vel_z * np.cos(theta)

        # Lorentz factor
        gamma = 1.0 / np.sqrt(1 - mag_v**2)

        # spin_phase
        spin_factor = np.exp(-2 * 1j * spin_phase)

        # Finally, the conformal factor
        conformal_factor = spin_factor * np.power(gamma * (1 - v_dot_r), 3)

        return conformal_factor

    def compute_supertransl_alpha(supertransl_alpha_modes, theta, phi):
        """Compute the spherical Alpha supertranslation variable
        :math:`\\alpha(\\theta, \\phi)` given its modes. This method
        just multiplies the alpha modes with their corresponding spherical
        harmonic basis functions and returns the summed result.

        Parameters
        ----------

        supertransl_alpha_modes:	dict
                                    A dictionary of lists, each sublist
                                    containing the set of super-translation
                                    modes corresponding to a particular
                                    :math:`\\ell'.
        theta:		float
                                The polar angle :math:`\\theta'.
        phi:	float
                        The azimuthal angle :math:`\\phi'.

        Returns
        --------

        supertransl_alpha_sphere:		func
                                        A function on the sphere
                                        (arguments :math:`\\theta', math:`\\phi').

        """

        # For partial evaluation of functions
        # from functools import partial
        print(supertransl_alpha_modes.keys())
        # Find the extreme ell values.
        keys_list = sorted(list(supertransl_alpha_modes.keys()))

        # ell_min = int(keys_list[0][1])
        # ell_max = int(keys_list[-1][1])

        # Import the Spherical Harmonic function

        from waveformtools.transforms import Yslm

        spin_weight = 0
        # Ylm = partial(Yslm, spin_weight=0)
        # The final function
        supertransl_alpha_sphere = 0

        theta = np.pi / 2
        phi = 0.0
        for item in keys_list:
            ell = int(item[1])
            for m_index in range(2 * ell + 1):
                emm = m_index - ell
                print("ell is", ell, type(ell), "emm is ", emm)
                supertransl_alpha_sphere += supertransl_alpha_modes[item][m_index] * Yslm(spin_weight, ell, emm, theta, phi)

        return supertransl_alpha_sphere

    def boost_waveform(unboosted_waveform, conformal_factor):
        """Boost the waveform given the unboosted waveform and the boost conformal factor.

        Parameters
        ----------

        non_boosted_waveform:	list
                                A list with a single floating point number
                                or a numpy array of the unboosted waveform.
                                The waveform can have angular as well as
                                time dimentions.

                                The nesting order should be that, given the
                                list `non_boosted_waveform', each item in the
                                list refers to an array defined on the sphere
                                at a particular time or frequency. The subitem
                                will have dimensions [ntheta, nphi].

        conformal_factor:		float/array
                                The conformal factor for the Lorentz transformation.
                                It may be a single floating point number or an array
                                on a spherical grid. The array will be of dimensions
                                [ntheta, nphi]

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

        """

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

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

        # Compute the meshgrid for theta and phi.
        # theta, phi = gridinfo.meshgrid

        # A list to store the boosted waveform.
        boosted_waveform = []

        for item in unboosted_waveform:
            # Compute the boosted waveform on the spherical grid on all the elements.

            # conformal_k_on_sphere = compute_conformal_k(vec_v, theta, phi)
            boosted_waveform_item = conformal_factor * item

            boosted_waveform.append(boosted_waveform_item)

        return boosted_waveform