# Source code for waveformtools.extrapolate

    """ Methods for waveform extrapolation. """

    ##############################################
    # Imports
    ##############################################
    import numpy as np

    import waveformtools

    ###############################################
    # Basic utilities
    ###############################################

    def r_to_ra_conversion(coord_radius, mass=1, spin=0):
        """Convert the isotropic co-ordinate radius parameter `r` in the ETK simulations
                into the approximate areal radius.

        Parameters
        ----------

        coord_radius:	float
                                        The coordinate radius in the Einstein toolkit

        mass:	float, optional
                        The sum of the quasi-local horizon (Christodolou) masses of the black holes.
                        Defaults to 1.

        spin:	float, optional
                        The magnitude of the spin of the system, as approximated by a single Kerr black hole
                        far away from the system. Defaults to 0.

        Returns
        -------

        areal_radius:	float
                                        The appriximate areal radius of the sphere.

        Notes
        -----

        Assumes that the system interoir to the sphere at co-ordinate radius `r_coord` is well approximated by a
        Kerr black hole.

        References
        ----------

        Nakano et al., (2015),	Phys. Rev. D 91, 104022, in-text below Eq.[30].
        """

        areal_radius = coord_radius * (1 + (mass + spin) / (2 * coord_radius)) * (1 + (mass - spin) / (2 * coord_radius))

        return areal_radius

    ############################################################################
    # Perturbative extraction
    ############################################################################

    def waveextract_to_inf_perturbative_twop5_order(rPsi4_rlm, delta_t, areal_radius=500, mass=1, spin=0, ell=2, emm=2):
        """Extract a numerical waveform to null infinity using perturbative techniques. This is :
                * accurate to second order in :math:`1/r`.
                * accurate to first order in Kerr mass and spin.
                * corrects for spheroidal harmonics

        Parameters
        ----------

        rPsi4_rlm:	1d array
                                The extracted Weyl scalar :math:`r\\Psi_{4\\ell m}` data array

        delta_t:	float
                                The time stepping.

        areal_radius:	1d array
                                        The areal radius of the extraction sphere.

        mass:	float
                        The total horizon mass of the system.

        spin:	float, optional
                        The effective spin of the spacetime. Defaults to 0.

        ell:	int
                        The polar quantum number :math:`\\ell`.

        emm:	int
                        The azimuthal quantum number :math:`m`.

        Returns
        -------

        rPsi4_inflm:	1d array
                                The waveform extracted to null infninity :math:`\\mathcal{I}^+`

        References
        ----------

        This implements the definition in Nakano et al., (2015),  Phys. Rev. D 91, 104022 Eq.[29].

        """

        from waveformtools.integrate import fixed_frequency_integrator
        from waveformtools.differentiate import differentiate_cwaveform

        # Timeaxis

        timeaxis = np.arange(0, len(rPsi4_rlm) * delta_t, delta_t)
        # Assigning the terms. Each set of subterms in a pair of paranthesis is a term.
        term_1_prefac = 1 - 2 * mass / areal_radius
        subterm_1_1 = rPsi4_rlm
        subterm_1_2_prefac = (ell - 1) * (ell + 2) / (2 * areal_radius)
        subterm_1_2, _ = fixed_frequency_integrator(rPsi4_rlm, delta_t, omega0=0.015)  # Integral_rPsi4_rlm
        subterm_1_3_prefac = (ell - 1) * (ell + 2) * (ell**2 + ell - 4) / (8 * areal_radius**2)
        subterm_1_3, _ = fixed_frequency_integrator(rPsi4_rlm, delta_t, order=2, omega0=0.015)  # Double_integral_rPsi4_rlm

        term_1 = term_1_prefac * (subterm_1_1 + subterm_1_2_prefac * subterm_1_2 + subterm_1_3_prefac * subterm_1_3)

        term_2_prefac = (2 * 1j * spin / (ell + 1) ** 2) * np.sqrt(
            (ell + 3) * (ell - 1) * (ell + emm + 1) * (ell - emm + 1) / ((2 * ell + 1) * (2 * ell + 3))
        )
        subterm_2_1 = differentiate_cwaveform(timeaxis, rPsi4_rlm)  # Differential of waveform
        subterm_2_2_prefac = -ell * (ell + 3) / areal_radius
        subterm_2_2 = subterm_1_1

        term_2 = term_2_prefac * (subterm_2_1 + subterm_2_2_prefac * subterm_2_2)

        term_3_prefac = (-2 * 1j * spin / ell**2) * np.sqrt(
            ((ell + 2) * (ell - 2) * (ell + emm) * (ell - emm)) / ((2 * ell - 1) * (2 * ell + 1))
        )
        subterm_3_1 = subterm_2_1
        subterm_3_2_prefac = -(ell - 2) * (ell + 1) / areal_radius
        subterm_3_2 = subterm_1_1

        term_3 = term_3_prefac * (subterm_3_1 + subterm_3_2_prefac * subterm_3_2)

        rPsi4_inflm = term_1 + term_2 + term_3

        return rPsi4_inflm

    def waveextract_to_inf_perturbative_two_order(rPsi4_rlm, delta_t, areal_radius=500, mass=1, ell=2):
        """Extract a numerical waveform to null infinity using perturbative techniques. This is :
                * accurate to second order in :math:`1/r`.
                * accurate to first order in Kerr mass and spin.
                * corrects for spheroidal harmonics

        Parameters
        ----------

        rPsi4_rlm:	1d array
                                The extracted Weyl scalar :math:`r\\Psi_{4\\ell m}` data array.

        delta_t:	float
                                The time stepping.

        areal_radius:	1d array
                                        The areal radius of the extraction sphere.

        mass:	float
                        The total horizon mass of the system.

        spin:	float, optional
                        The effective spin of the spacetime. Defaults to 0.

        ell:	int
                        The polar quantum number :math:`\\ell`.

        emm:	int
                        The azimuthal quantum number :math:`m`.

        Returns
        -------

        rPsi4_inflm:	1d array
                                The waveform extracted to null infninity :math:`\\mathcal{I}^+`

        References
        ----------

        This implements the definition in Nakano et al., (2015),  Phys. Rev. D 91, 104022 Eq.[29].

        """

        from integrate import fixed_frequency_integrator

        # Assigning the terms. Each set of subterms in a pair of paranthesis is a term.
        term_1 = rPsi4_rlm

        term_2_prefac = -(ell - 1) * (ell + 2) / (2 * areal_radius)
        subterm_2_1, _ = fixed_frequency_integrator(rPsi4_rlm, delta_t, omega0=0.015)  # Integral_rPsi4_rlm
        term_2 = term_2_prefac * subterm_2_1

        term_3_prefac = (ell - 1) * (ell + 2) * (ell**2 + ell - 4) / (8 * areal_radius**2)
        subterm_3_1, _ = fixed_frequency_integrator(rPsi4_rlm, delta_t, order=2, omega0=0.015)  # Double_Integral_rPsi4_rlm
        term_3 = term_3_prefac * subterm_3_1

        term_4_prefac = -3 * mass / (2 * areal_radius**2)
        subterm_4_1 = subterm_2_1
        term_4 = term_4_prefac * subterm_4_1

        rPsi4_inflm = term_1 + term_2 + term_3 + term_4

        return rPsi4_inflm

    def waveextract_to_inf_perturbative_one_order(u_ret, rPsi4_rlm, areal_radius=500, ell=2):
        """Extract a numerical waveform to null infinity using perturbative techniques. This is :
                * accurate to second order in :math:`1/r`.
                * accurate to first order in Kerr mass and spin.
                * corrects for spheroidal harmonics

        Parameters
        ----------

        u_ret:	1d array
                        The retarted time array at the location r = areal_radius.

        rPsi4_rlm:	1d array
                                The extracted Weyl scalar :math:`r\\Psi_{4\\ell m}` data array

        areal_radius:	1d array
                                        The areal radius of the extraction sphere.

        mass:	float
                        The total horizon mass of the system.

        ell:	int
                        The polar quantum number :math:`\\ell`.

        emm:	int
                        The azimuthal quantum number :math:`m`.

        Returns
        -------

        rPsi4_inflm:	1d array
                                The waveform extracted to null infninity :math:`\\mathcal{I}^+`

        References
        ----------

        This implements the definition in Nakano et al., (2015),  Phys. Rev. D 91, 104022 Eq.[29].

        """

        # Get the amplitude and phase
        A_lm, _ = waveformtools.xtract_camp_phase(rPsi4_rlm.real, rPsi4_rlm.imag)

        # Get the time stepping
        delta_t = u_ret[1] - u_ret[0]

        # Get the waveform instantaneous frequency
        omega_lm = waveformtools.get_waveform_angular_frequency(rPsi4_rlm, delta_t=delta_t, timeaxis=u_ret)

        # Assigning the terms. Each set of subterms in a pair of paranthesis is a term.

        # The amplitude correction factor
        A_lm_correction = 0.5 * np.power(((ell * (ell + 1) / (2 * omega_lm * areal_radius))), 2)

        # The phase correction factor.
        sin_Phase_correction = ell * (ell + 1) / (2 * omega_lm * areal_radius)
        cos_Phase_correction = np.sqrt(1 - np.power(sin_Phase_correction, 2))
        Phase_correction = cos_Phase_correction + 1j * sin_Phase_correction

        # The extrapolated waveform

        A_lm_corrected = A_lm * (1 + A_lm_correction)
        Phi_lm_corrected = np.exp(-1j * omega_lm * u_ret) * Phase_correction

        rPsi4_inflm = A_lm_corrected * Phi_lm_corrected

        return rPsi4_inflm