"""
JaxGB parameter module
======================

This is a module clarifying the parameters (names, units, and LaTeX notation) and
their ordering in JaxGB. It also enables conversions to/from the parameterisation
used in other tools, by defining a `GBObject` containing all necessary parameters.

.. autoclass:: GBObject
    :members:
"""

import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from astropy.coordinates import SkyCoord

PARAM_NAMES = np.array(
    [
        "Frequency",
        "Frequency derivative",
        "Amplitude",
        "Right ascension",
        "Declination",
        "Polarisation angle",
        "Inclination",
        "Initial phase",
    ]
)
"""Array of ordered parameter names."""

PARAM_LATEX = np.array(
    [
        "$f_0$",
        r"$\dot{f}$",
        "$A$",
        r"$\mathrm{RA}$",
        r"$\mathrm{dec}$",
        r"$\psi$",
        r"$\iota$",
        r"$\phi_0$",
    ]
)
"""Array of ordered Latex-formatted parameter symbols."""

PARAM_UNITS = np.array(
    [
        r"$\mathrm{Hz}$",
        r"$\mathrm{Hz}/\mathrm{s}$",
        "",
        r"$\mathrm{rad}$",
        r"$\mathrm{rad}$",
        r"$\mathrm{rad}$",
        r"$\mathrm{rad}$",
        r"$\mathrm{rad}$",
    ]
)
"""Array of ordered parameter expected units."""


class GBObject:
    """Circular Galactic binary parameter set.

    GBObject brings together a canonical parameterisation of a set of circular
    Galactic Binaries. It enables conversion into the parameterisation used by
    different (circular) Galactic Binary response codes, like LISA GW Response, LDC
    tools, fastgb, and gbgpu. At the moment, we ignore the inclusion of a second
    derivative of the frequency.

    Parameters
    ----------
    f0
        Gravitational-wave signal frequency at epoch [Hz].
    fdot
        Gravitational-wave signal frequency derivative [Hz/s].
    A
        Gravitational-wave signal amplitude [strain].
    ra
        Right ascension (equatorial longitude) [rad].
    dec
        Declination, equatorial latitude [rad].
    psi
        Polarisation angle [rad].
    iota
        Inclination angle [rad].
    phi0
        Initial phase at epoch [rad].
    t_init
        Epoch used to define the initial phase phi0 and the frequency f0, in TCB [s].

    """

    def __init__(
        self,
        *,
        f0: np.ndarray,
        fdot: np.ndarray,
        A: np.ndarray,
        ra: np.ndarray,
        dec: np.ndarray,
        psi: np.ndarray,
        iota: np.ndarray,
        phi0: np.ndarray,
        t_init: float,
    ) -> None:

        self.f0 = f0
        """Gravitational-wave signal frequency at epoch [Hz]."""
        self.fdot = fdot
        """Gravitational-wave signal frequency derivative [Hz/s]."""
        self.A = A
        """Gravitational-wave signal amplitude [strain]."""
        self.ra = ra
        """Right ascension (equatorial longitude) [rad]."""
        self.dec = dec
        """Declination, equatorial latitude [rad]."""
        self.psi = psi
        """Polarisation angle [rad]."""
        self.iota = iota
        """Inclination angle [rad]."""
        self.phi0 = phi0
        """Initial phase at epoch [rad]."""
        self.t_init = t_init
        """Epoch used to define the initial phase phi0 and the frequency f0, in TCB
        [s]."""

    @classmethod
    def from_jaxgb_params(cls, jaxgb_params: np.ndarray, t_init: float) -> "GBObject":
        """Instantiate GB from a JaxGB params array.

        Parameters
        ----------
        jaxgb_params : array-like of shape (nsrc, 8)
            Parameters, as defined in :attr:`PARAM_NAMES` and units
            as defined in :attr:`PARAM_UNITS`.
        t_init : float
            Epoch used to define the initial phase phi0 and the frequency f0, in TCB
            [s].

        Returns
        -------
            A new GB instance.
        """
        gb_object = cls(
            f0=jaxgb_params[:, 0],
            fdot=jaxgb_params[:, 1],
            A=jaxgb_params[:, 2],
            ra=jaxgb_params[:, 3],
            dec=jaxgb_params[:, 4],
            psi=jaxgb_params[:, 5],
            iota=jaxgb_params[:, 6],
            phi0=jaxgb_params[:, 7],
            t_init=t_init,
        )
        return gb_object

    @classmethod
    def from_pandas_dataframe(
        cls, pandas_df: pd.DataFrame, t_init: float
    ) -> "GBObject":
        """Instantiate GB from a Pandas dataframe.

        Parameters
        ----------
        pandas_df
            Pandas dataframe of parameters.
        t_init : float
            Epoch used to define the initial phase phi0 and the frequency f0, in TCB
            [s].

        Returns
        -------
            A new GB instance.
        """
        gb_object = cls(
            f0=pandas_df["Frequency"].to_numpy(),
            fdot=pandas_df["FrequencyDerivative"].to_numpy(),
            A=pandas_df["Amplitude"].to_numpy(),
            ra=pandas_df["RightAscension"].to_numpy(),
            dec=pandas_df["Declination"].to_numpy(),
            psi=pandas_df["Polarization"].to_numpy(),
            iota=pandas_df["Inclination"].to_numpy(),
            phi0=pandas_df["InitialPhase"].to_numpy(),
            t_init=t_init,
        )
        return gb_object

    @classmethod
    def from_records(cls, rec_array: np.recarray | dict, t_init: float) -> "GBObject":
        """Instantiate GB from a numpy record array.

        Parameters
        ----------
        rec_array
            Numpy record array or dictionnary of parameters.
        t_init : float
            Epoch used to define the initial phase phi0 and the frequency f0, in TCB
            [s].

        Returns
        -------
            A new GB instance.
        """
        gb_object = cls(
            f0=rec_array["Frequency"],
            fdot=rec_array["FrequencyDerivative"],
            A=rec_array["Amplitude"],
            ra=rec_array["RightAscension"],
            dec=rec_array["Declination"],
            psi=rec_array["Polarization"],
            iota=rec_array["Inclination"],
            phi0=rec_array["InitialPhase"],
            t_init=t_init,
        )
        return gb_object

    @staticmethod
    def _equatorial_to_ecliptic(ra: float, dec: float) -> tuple[float, float]:
        """Convert (ra, dec) into ecliptic longitude and lattitude (lambda, beta).

        Parameters
        ----------
        ra
            Right ascension [rad].
        dec
            Declination [rad].

        Returns
        -------
        tuple (lambda, beta)
            Ecliptic longitude and latitude [rad, rad].
        """
        sky_icrs = SkyCoord(ra=ra, dec=dec, frame="icrs", unit="rad")
        ecl_coords = sky_icrs.barycentrictrueecliptic
        return ecl_coords.lon.rad, ecl_coords.lat.rad

    def _to_array_at_reference_time(self, t0: float) -> np.ndarray:
        """Create a numpy array with parameter at reference time `t0`.

        Returns
        -------
        np.array of shape (nsrc, 8)
            Array of parameters in the same order as :attr:`PARAM_NAMES`
            and with units :attr:`PARAM_UNITS`.
        """
        # shift frequency `f0`, which evolves at the linear rate `fdot`
        f0_t0 = self.f0 - self.fdot * (self.t_init - t0)

        # shift `phi0`, to account for difference `t_init` (catalogue epoch) and `t0`
        _shift = np.pi * (
            2 * self.f0 * (self.t_init - t0) - self.fdot * (self.t_init - t0) ** 2
        )
        phi0_t0 = (self.phi0 + _shift) % (2 * np.pi)

        return np.array(
            [
                f0_t0,  # shifted f0
                self.fdot,
                self.A,
                self.ra,
                self.dec,
                self.psi,
                self.iota,
                phi0_t0,  # shifted phi0
            ]
        ).T

    def to_jaxgb_array(self, t0: float) -> jax.Array:
        """Create JaxGB array, respecting format and ordering.

        Parameters
        ----------
        t0
            Starting time of `JAXGB` response calculation in TCB (s). This time is used
            to shift the GW frequency `f0` and the phase `phi0`, defined at `t_init`, to
            the simulation starting time.

        Returns
        -------
        jax.Array of shape (nsrc, 8)
            Array of parameters in the same order as :attr:`PARAM_NAMES`
            and with units :attr:`PARAM_UNITS`.
        """
        return jnp.array(self._to_array_at_reference_time(t0=t0))

    def to_ldc_dict(self) -> dict[str, np.ndarray]:
        """Create LDC dictionary, respecting format and ordering.

        Sky positions are converted to ecliptic coordinates. Parameters are specified
        at the epoch `t_init`.

        Returns
        -------
            Dictionary of parameters.
        """
        # LDC assumes no difference between `t0` and `t_init`
        params = self._to_array_at_reference_time(t0=self.t_init)

        # convert equatorial to ecliptic coordinates, change order
        params[4], params[3] = GBObject._equatorial_to_ecliptic(params[3], params[4])

        return dict(
            zip(
                [
                    "Frequency",
                    "FrequencyDerivative",
                    "Amplitude",
                    "EclipticLatitude",
                    "EclipticLongitude",
                    "Polarization",
                    "Inclination",
                    "InitialPhase",
                ],
                params,
            )
        )

    def to_fastgb_array(self) -> np.ndarray:
        """Create fastgb array, respecting format and ordering.

        The parameter array assumes the epoch `t_init`, and fastgb is doing the
        conversion between `tinit` and `t0` internally.

        Additional comments:

        fastgb expects that `params[:, 3]` contain ecliptic latitudes (beta) and
        `params[:, 4]` contain ecliptic longitudes (lambda). However, from LISA
        Orbits v3 onwards, orbit information is expressed in equatorial coordinates,
        so implictly, the sky location params in fastgb need to be expressed in
        equatorial coordinates. In principle, we can therefore pass the equatorial
        coordinates, as in JAXGB, provided we swap the longitude and latitude order:

        - In JAXGB, `params[:, 3]` is longitude (ra) while `params[:, 4]` includes
          latitude (dec),
        - In fastgb, `params[:, 3]` (what used to be beta) includes latitude, while
          `params[:, 4]` includes longitude (what used to be beta)

        Returns
        -------
        np.ndarray of shape (nsrc, 8)
            Array of parameters.
        """
        # fastgb applies internal conversion between `t0` and `t_init`, so we assume
        # them to be identical here
        params = self._to_array_at_reference_time(t0=self.t_init)
        return params[[0, 1, 2, 4, 3, 5, 6, 7]]

    def to_gwresponse_dict(self) -> dict[str, np.ndarray]:
        """Create GW Response input dictionary, respecting format and ordering.

        This can be used with :class:`lisagwresponse.GalacticBinary`. The parameter
        dictionary assumes the epoch `t_init`.

        Returns
        -------
            Dictionary of parameters, compatible with
            :class:`lisagwresponse.GalacticBinary`.
        """
        # GW response applies own internal time shift for `t0` and `t_init`, so we
        # assume them to be identical here
        params = self._to_array_at_reference_time(t0=self.t_init)
        return dict(
            zip(
                [
                    "f",
                    "df",
                    "A",
                    "ra",
                    "dec",
                    "psi",
                    "iota",
                    "phi0",
                ],
                params,
            )
        )

    def to_gbgpu_array(self) -> np.ndarray:
        """Create gbgpu array, respecting format and ordering.

        Sky positions are converted to ecliptic coordinates. ``fddot`` is
        assumed to be zero. The returned parameters have the epoch `t_init`.

        Returns
        -------
        np.ndarray of shape (nsrc, 9)
            Array of parameters.
        """
        # gbgbpu assumes no difference between `t0` and `t_init`
        params = self._to_array_at_reference_time(t0=self.t_init)
        params = params[[2, 0, 1, 1, 7, 6, 5, 3, 4]]
        params[3] = 0  # f2dot
        # convert equatorial to corresponding ecliptic coordinates
        params[7], params[8] = GBObject._equatorial_to_ecliptic(params[7], params[8])
        return np.array([np.full(1, p) for p in params])
