"""
JaxGB TDI module
================

This is a placeholder waiting for a third-party software managing TDI
combinations and generations.

For now it gives basic transformations from XYZ (Michelson) to AET
(orthogonal) assuming equal diagonal and cross terms in the noise matrix,
and basic transformation from TDI 1.5 to TDI 2 assuming equal
armlengths.

.. autofunction:: to_tdi_generation
.. autofunction:: to_tdi_combination
"""

import jax
import jax.numpy as jnp
import numpy as np
from lisaconstants import c  # pylint: disable=no-name-in-module


def to_tdi_generation(
    tdi: jax.Array,
    f: jax.Array,
    to: float = 1.5,
    arm_length: jax.Array = jnp.array(2.5e9),
) -> jax.Array:
    """Frequency-domain transformation of TDI 1.5 to another generation.

    The transformation is done in the frequency domain, assuming equal
    armlengths, ie., a single factor.

    Parameters
    ----------
    tdi
        Input TDI in the frequency domain of shape
        (nsource, 3, nfreq).
    f
        Frequency of shape nfreq [Hz]
    to
        either 1.5 or 2
    arm_length
        Constant armlengths [m].

    Returns
    -------
    Array of shape (nsource, 3, freq)
        Output TDI in the frequency domain.
    """
    if to == 2:
        omegal = 2 * np.pi * f * jnp.mean(arm_length) / c
        fctr = 2.0j * jnp.sin(2 * omegal) * jnp.exp(-2j * omegal)
        return tdi * fctr[:, None, :]
    return tdi


def to_tdi_combination(tdi: jax.Array, combination: str = "XYZ") -> jax.Array:
    """Transform input TDI XYZ combination into given combination.

    This can be applied to time or frequency series.

    Parameters
    ----------
    tdi
        Input TDI of shape (nsource, 3, nfreq).
    combination
        Combination name, can be XYZ or AET.

    Returns
    -------
    Array of shape (nsource, 3, freq)
        Output TDI.
    """
    if combination == "AET":
        X, Y, Z = tdi[:, 0, :], tdi[:, 1, :], tdi[:, 2, :]
        A, E, T = (
            (Z - X) / np.sqrt(2.0),
            (X - 2.0 * Y + Z) / np.sqrt(6.0),
            (X + Y + Z) / np.sqrt(3.0),
        )
        tdi = jnp.stack([A, E, T], axis=1)
        return tdi
    return tdi
