"""Benchmark egw computation using full vs small segment"""

from gw_eccentricity.load_data import load_waveform
from gw_eccentricity import measure_eccentricity
from copy import deepcopy

lal_kwargs_10_orbits = {
    "approximant": "EccentricTD",
    "q": 1.0,
    "chi1": [0.0, 0.0, 0.0],
    "chi2": [0.0, 0.0, 0.0],
    "Momega0": 0.0155,
    "ecc": 0.1,
    "mean_ano": 0,
    "include_zero_ecc": True}
lal_kwargs_100_orbits = deepcopy(lal_kwargs_10_orbits)
lal_kwargs_100_orbits["Momega0"] = 0.00524
# Make a dataDict for waveforms with different number of orbits
dataDict_10_orbits = load_waveform(**lal_kwargs_10_orbits)
dataDict_100_orbits = load_waveform(**lal_kwargs_100_orbits)
# reference times
trefs = {"10": -3960,
         "100": -89602} 

def compute_eccentricity(tref, data_dict, use_segment):
    measure_eccentricity(
        tref_in=tref,
        dataDict=data_dict,
        method="AmplitudeFits",
        extra_kwargs={
            "use_segment": use_segment,
        })

def test_short_vs_full_use_segment_True_10_orbits(benchmark):
    benchmark(compute_eccentricity, trefs["10"], dataDict_10_orbits, True)

def test_short_vs_full_use_segment_False_10_orbits(benchmark):
    benchmark(compute_eccentricity, trefs["10"], dataDict_10_orbits, False)

def test_short_vs_full_use_segment_True_100_orbits(benchmark):
    benchmark(compute_eccentricity, trefs["100"], dataDict_100_orbits, True)

def test_short_vs_full_use_segment_False_100_orbits(benchmark):
    benchmark(compute_eccentricity, trefs["100"], dataDict_100_orbits, False)
