"""Compare JAXGB and GBGPU fast waveforms."""

import jax
import jax.numpy as jnp
import lisaorbits
import numpy as np
from fastgb import fastgb
from gbgpu.gbgpu import GBGPU
from testutils import get_source, overlap

from jaxgb.jaxgb import JaxGB
from jaxgb.tdi import to_tdi_combination

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


def test_gbgpu(gw_name: str, tdi_generation: float) -> None:
    """Check that overlap between JaxGB and GBGPU

    We assume equal armlength orbits.
    """
    size, dt, tmax = 2048, 15, 180 * 24 * 3600

    # GBGPU
    gb = GBGPU(force_backend="cpu")
    params = get_source(gw_name, as_gbgpu=True)
    gb.run_wave(*params, N=size, dt=dt, T=tmax, oversample=1, tdi2=tdi_generation == 2)

    A1 = gb.A[0]
    E1 = gb.E[0]
    freqs = gb.freqs[0]

    fgb = JaxGB(orbits=lisaorbits.EqualArmlengthOrbits(), t_obs=tmax, t0=0, n=size)
    params = get_source(gw_name)
    A2, E2, _ = fgb.get_tdi(
        params,
        tdi_generation=tdi_generation,
        tdi_combination="AET",
    )
    kmin = fgb.get_kmin(params[0])  # frequency

    assert A1.shape == A2.shape
    df = 1 / tmax
    assert jnp.isclose(freqs, df * jnp.arange(kmin, kmin + A2.shape[0])).all()
    olap = overlap(np.hstack([A1, E1]), np.hstack([A2, E2]))
    assert olap > 0.965


def test_fastgb(gw_name: str, tdi_generation: float) -> None:
    """Check that overlap between JaxGB and fastGB

    We assume equal armlength orbits.
    """
    size, dt, tmax = 2048, 15, 180 * 24 * 3600

    # fastGB
    fgb = fastgb.FastGB(
        delta_t=dt, T=tmax, N=size, orbits=lisaorbits.EqualArmlengthOrbits()
    )
    params = get_source(gw_name, as_fastgb=True)
    X, Y, Z, _ = fgb.get_fd_tdixyz(
        jnp.array(params).reshape(1, -1), tdi2=tdi_generation == 2
    )
    tdi = to_tdi_combination(jnp.stack([X, Y, Z], axis=1), "AET")
    A1, E1 = tdi[:, 0, :], tdi[:, 1, :]

    # JAXGB
    params = get_source(gw_name)
    jgb = JaxGB(orbits=lisaorbits.EqualArmlengthOrbits(), t_obs=tmax, t0=0, n=size)
    A2, E2, _ = jgb.get_tdi(
        params,
        tdi_generation=tdi_generation,
        tdi_combination="AET",
    )

    olap = overlap(np.hstack([A1, E1]), np.hstack([A2, E2]))
    assert olap > 0.9999
