"""
Test TDI computation of several sources at one time, either with
jax.vmap or multi columns parameters
"""

import jax
import jax.numpy as jnp
import lisaorbits
import numpy as np
import pandas as pd

from jaxgb.jaxgb import JaxGB
from jaxgb.params import GBObject

jax.config.update("jax_enable_x64", True)


def test_multicol(catalog: pd.DataFrame, tdi_generation: float) -> None:
    """Test the computation of tdi for several sources at one time, using
    2D input array.
    """

    size, tmax = 512, 180 * 24 * 3600
    ncol = len(catalog)
    fgb = JaxGB(orbits=lisaorbits.EqualArmlengthOrbits(), t_obs=tmax, t0=0, n=size)

    # create JAXGB compatible input array from catalog
    params = GBObject.from_pandas_dataframe(catalog, t_init=0.0).to_jaxgb_array(t0=0.0)

    X, _, _ = fgb.get_tdi(params, tdi_generation=tdi_generation)
    assert X.shape == (ncol, size)

    i_check = 10
    x, _, _ = fgb.get_tdi(params[i_check], tdi_generation=tdi_generation)
    assert np.isclose(x, X[i_check], atol=1e-25).all()


def test_vmap(catalog: pd.DataFrame) -> None:
    """Test the computation of tdi for several sources at one time, using
    jax.vmap
    """

    size, tmax = 512, 180 * 24 * 3600
    ncol = len(catalog)
    fgb = JaxGB(orbits=lisaorbits.EqualArmlengthOrbits(), t_obs=tmax, t0=0, n=size)

    # create JAXGB compatible input array from catalog
    params = GBObject.from_pandas_dataframe(catalog, t_init=0.0).to_jaxgb_array(t0=0.0)

    X1, _, _ = jax.vmap(fgb.get_tdi, in_axes=0)(params)
    assert X1.shape == (ncol, size)

    X2, _, _ = fgb.get_tdi(params)
    assert np.isclose(X1, X2, atol=1e-25).all()


def test_sum(catalog: pd.DataFrame) -> None:
    """Test the computation of the summed tdi for several sources at one time"""

    size, tmax = 512, 180 * 24 * 3600

    # create JAXGB compatible input array from catalog
    params = GBObject.from_pandas_dataframe(catalog, t_init=0.0).to_jaxgb_array(t0=0.0)

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

    # test for one source
    i_check = 10
    x, _, _ = fgb.get_tdi(params[i_check])
    X, _, _ = fgb.sum_tdi(jnp.vstack([params[i_check], params[i_check]]))
    assert np.isclose(X, 2 * x, atol=1e-25).all()

    X1, _, _ = fgb.sum_tdi(params, kmin=0, kmax=100_000)
    assert X1.shape == (100_000,)

    # assuming no overlap of the highest freq source with any other
    catalog = catalog.sort_values(by=["Frequency"])
    params = GBObject.from_pandas_dataframe(catalog, t_init=0.0).to_jaxgb_array(t0=0.0)
    last = params[len(catalog) - 1, 0]
    x, _, _ = fgb.get_tdi(params[len(catalog) - 1])  # select last source
    kmin = fgb.get_kmin(last)  # params[i_last, 0] = f0
    X, _, _ = fgb.sum_tdi(params, kmin=kmin, kmax=kmin + size)
    assert fgb.get_kmin(params[len(catalog) - 2, 0]) + size < kmin
    assert np.isclose(X, x, atol=1e-25).all()
