"""Check that JAXGB waveform can be auto-differentiated by JAX."""

import jax
import jax.numpy as jnp
import lisaorbits
from testutils import get_source
from wejax.jacobian import autograd

from jaxgb.jaxgb import JaxGB

# pylint: disable=too-many-locals
jax.config.update("jax_enable_x64", True)


def test_wejax(gw_name: str) -> None:
    """Check that autograd is simply working"""
    size, _, tmax = 2048, 15, 180 * 24 * 3600

    fgb = JaxGB(orbits=lisaorbits.EqualArmlengthOrbits(), t_obs=tmax, t0=0, n=size)

    def tdi(*pars):
        """
        Generate the template frequency serie for the TDI channels
        (A, E).

        Parameters
        ----------
        pars
            template parameters.

        Returns
        -------
           jax.ndarray (Nf x Nc) with the A and E signals (we ignore the T channel)
        """
        A, E, _ = fgb.get_tdi(jnp.array(pars), tdi_combination="AET")
        return jnp.stack([A, E], axis=1)

    params = get_source(gw_name)
    grad_gen = autograd(tdi, argnums=list(range(8)))
    gb_grad = grad_gen(*params)
    assert gb_grad.shape == (size, 2, 8)
