"""Compare fast waveform to time domain generated waveform using
gwresponse and pytdi.

The comparison is performed in the frequency domain.
"""

import os

import h5py
import jax
import jax.numpy as jnp
import lisagwresponse
import lisaorbits
import numpy as np
import pandas as pd
import pytdi.michelson
import pytest
from jax import Array
from lisaconstants.indexing import LINKS
from lisaorbits import ResampledOrbits
from pytdi import Data
from testutils import get_source, mismatch, overlap

from jaxgb.jaxgb import JaxGB, JaxGBaccurate
from jaxgb.params import GBObject
from jaxgb.tdi import to_tdi_combination

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

# pylint: disable=too-many-locals
# pylint: disable=redefined-outer-name


def _read_t0(orbits_file):
    with h5py.File(orbits_file, "r") as fid:
        t0 = fid.attrs["t0"]
    return t0


def _gw2tdi(
    orbits_fn: str,
    gwfile: str,
    dt: int,
    tdi_generation: float = 2,
    tdi_combination: str = "XYZ",
) -> tuple[Array, Array, Array, Array]:
    """Return TDI X,Y,Z channel and associated frequencies from link
    response stored in `gwfile`.
    """
    data = Data.from_gws(gwfile, orbits_fn)
    if tdi_generation == 2:
        X = pytdi.michelson.X2.build(**data.args)(data.measurements)
        Y = pytdi.michelson.Y2.build(**data.args)(data.measurements)
        Z = pytdi.michelson.Z2.build(**data.args)(data.measurements)
    else:
        X = pytdi.michelson.X1.build(**data.args)(data.measurements)
        Y = pytdi.michelson.Y1.build(**data.args)(data.measurements)
        Z = pytdi.michelson.Z1.build(**data.args)(data.measurements)

    ## Move to freq domain
    freqs = np.fft.rfftfreq(X.size, d=dt)
    X = np.fft.rfft(X) * dt
    Y = np.fft.rfft(Y) * dt
    Z = np.fft.rfft(Z) * dt
    tdi = to_tdi_combination(
        jnp.stack([X.reshape(1, -1), Y.reshape(1, -1), Z.reshape(1, -1)], axis=1),
        tdi_combination,
    )
    return freqs, tdi[:, 0, :].squeeze(), tdi[:, 1, :].squeeze(), tdi[:, 2, :].squeeze()


def test_tdi_overlap(
    tmp_path: str | os.PathLike,
    gw_name: str,
    tdi_generation: str,
    orbits_approximation: str,
) -> None:
    """Check that overlap between fast waveform and time domain waveform
    is above 0.999 for a given source, orbits approximation and tdi
    generation.

    """
    size, dt, tmax = 2048, 15, 180 * 24 * 3600

    ## compute GW response and TDI
    gwfile = str(tmp_path / "gw.h5")
    p = get_source(gw_name, as_gwresponse_dict=True)
    p["orbits"] = orbits_approximation
    gb = lisagwresponse.GalacticBinary(**p)
    gb.write(gwfile, mode="w", t0=0, size=tmax // dt, dt=dt)
    freqs, A1, E1, _ = _gw2tdi(
        orbits_approximation,
        gwfile,
        dt,
        tdi_generation=tdi_generation,
        tdi_combination="AET",
    )

    ## compute fast approx
    fgb = JaxGB(orbits=ResampledOrbits(orbits_approximation), 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

    ## align on A/E
    df = 1 / tmax
    i0 = np.where(freqs == df * kmin)[0][0]
    sl = slice(i0, i0 + A2.shape[0])
    olap = overlap(np.hstack([A1[sl], E1[sl]]), np.hstack([A2, E2]))
    assert olap > 0.999


def test_link_overlap(gw_name: str, orbits_approximation: str) -> None:
    """Check that overlap at the link level, between fast waveform and
    time domain waveform is above 0.999 for a given source, orbits
    approximation and tdi generation.
    Each link is tested independently, to avoid mismatch cancellation.

    """
    size, dt, tmax = 2048, 15, 180 * 24 * 3600

    ## compute GW response
    p = get_source(gw_name, as_gwresponse_dict=True)
    p["orbits"] = orbits_approximation
    gb = lisagwresponse.GalacticBinary(**p)
    times = np.arange(0, tmax, dt).astype(float)
    yt = gb.compute_gw_response(times, LINKS)
    yf1 = np.fft.rfft(yt.T) * dt
    freqs = np.fft.rfftfreq(yt.shape[0], d=dt)

    ## compute fast response
    fgb = JaxGB(orbits=ResampledOrbits(orbits_approximation), t_obs=tmax, t0=0, n=size)

    params = get_source(gw_name)
    yf2 = fgb.get_link_responses(params.reshape(1, -1))[0, :, :]
    kmin = fgb.get_kmin(params[0])  # frequency

    ## align
    df = 1 / tmax
    i0 = np.where(freqs == df * kmin)[0][0]
    sl = slice(i0, i0 + yf2.shape[1])
    for link in range(6):
        olap = overlap(yf1[link, sl], np.array(yf2[link, :]))
        assert olap > 0.999


def test_esa_orbits(
    tmp_path: str | os.PathLike,
    gw_name: str,
    tdi_generation: str,
) -> None:
    """Check that overlap between fast waveform and time domain waveform
    is above 0.99 when using esa orbits. Handling t0!=0 case.
    """
    # esa orbits
    tmax = 730 * 24 * 3600
    orbits = lisaorbits.OEMOrbits.from_included("esa-leading")
    t0 = orbits.t_init + 50000000
    orbits_fn = str(tmp_path / "orbits.h5")
    orbits.write(orbits_fn, mode="w", dt=1e5, size=(2 * tmax) // 1e5, t0=t0 - 100)

    size, dt, tmax = 4096, 15, 180 * 24 * 3600

    ## compute GW response and TDI
    gwfile = str(tmp_path / "gw.h5")
    p = get_source(gw_name, as_gwresponse_dict=True)
    p["orbits"] = orbits_fn
    p["t_init"] = t0
    gb = lisagwresponse.GalacticBinary(**p)
    gb.write(gwfile, mode="w", t0=t0, size=tmax // dt, dt=dt)

    freqs, A1, E1, _ = _gw2tdi(
        orbits_fn, gwfile, dt, tdi_generation=tdi_generation, tdi_combination="AET"
    )

    ## compute fast approx
    fgb = JaxGB(orbits=orbits, t_obs=tmax, t0=t0, 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

    ## align on A/E
    df = 1 / tmax
    i0 = np.where(freqs == df * kmin)[0][0]
    sl = slice(i0, i0 + A2.shape[0])
    olap = overlap(np.hstack([A1[sl], E1[sl]]), np.hstack([A2, E2]))
    assert olap > 0.99


def test_accurate(gw_name: str, orbits_approximation: str) -> None:
    """Check overlap of accurate fast waveform versus time domain at the
    link level, in time domain.
    """
    size, dt, tmax = 4096, 15, 180 * 24 * 3600
    t0 = 0

    ## fast response
    fgb = JaxGBaccurate(
        orbits=ResampledOrbits(orbits_approximation),
        t_obs=tmax,
        t0=t0,
        n=size,
        window=0.5,
    )
    params = get_source(gw_name)
    yf = fgb.get_link_responses(params.reshape(1, -1))[0, :, :]
    kmin = fgb.get_kmin(params[0])  # frequency

    ## time domain
    p = get_source(gw_name, as_gwresponse_dict=True)
    p["orbits"] = orbits_approximation
    p["t_init"] = t0
    gb = lisagwresponse.GalacticBinary(**p)
    times = np.arange(t0, t0 + fgb.t_obs, dt).astype(float)
    yt = gb.compute_gw_response(times, LINKS)

    df = 1 / fgb.t_ext
    target_length = int(np.rint(1 / df / dt))

    for link in range(6):
        array = np.pad(yf[link, :], (kmin, 0))
        yt_fast = np.fft.irfft(array / dt, n=target_length)
        times_fast = fgb.tm[0] + np.arange(target_length) * dt

        # align
        i0 = np.where(times_fast == t0)[0][0]
        i1 = i0 + len(times)
        olap = overlap(yt[:, link], yt_fast[i0:i1])
        assert olap > 0.999
        mm = mismatch(yt[:, link], yt_fast[i0:i1])
        assert mm < 5e-15


@pytest.mark.parametrize("t_init", [97729930.202664, 50729089.327664])
def test_tinit(tmp_path: str | os.PathLike, t_init: float) -> None:
    """test t_init = t0 and t_init != t0

    a huge time difference is needed to be sensitive to the phase
    shift.
    """

    t0 = 97729930.202664

    src = pd.DataFrame(
        {
            "Amplitude": [4.184033623373269e-23],
            "RightAscension": [4.577818480741154],
            "Declination": [-0.4451836132523424],
            "Polarization": [2.1401202483554256],
            "Inclination": [2.5150470519327444],
            "FrequencyDerivative": [3.425955521564606e-16],
            "Frequency": [0.006094623731395194],
            "InitialPhase": [0.8363015305549726],
        }
    )
    gbo = GBObject.from_pandas_dataframe(src, t_init=t_init)
    params0 = gbo.to_jaxgb_array(t0=t_init)  # to be given to gw_response, no shift
    params = gbo.to_jaxgb_array(t0=t0)  # to be given to jaxgb, with shift

    duration = 24 * 3600 * 30 * 4
    orbits_sampling = {
        "dt": np.float64(500000.0),
        "duration": np.float64(137000000.0),
        "fs": np.float64(2e-06),
        "size": np.int64(274),
        "t0": np.float64(61171239.327664),
    }

    # esa orbits
    orbits_fn = str(tmp_path / "orbits.h5")
    # _orbits = lisaorbits.OEMOrbits.from_included("esa-trailing")
    _orbits = lisaorbits.EqualArmlengthOrbits()

    _orbits.write(
        orbits_fn,
        mode="w",
        dt=orbits_sampling["dt"],
        size=orbits_sampling["size"],
        t0=orbits_sampling["t0"],
    )

    size, dt = 1024, 15

    ## compute GW response and TDI
    gwfile = str(tmp_path / "gw.h5")
    p = dict(
        zip(
            [
                "f",
                "df",
                "A",
                "ra",
                "dec",
                "psi",
                "iota",
                "phi0",
            ],
            [float(p) for p in params0.squeeze()],
        )
    )
    p["orbits"] = orbits_fn
    p["t_init"] = t_init
    gb = lisagwresponse.GalacticBinary(**p)
    gb.write(gwfile, mode="w", t0=t0, size=duration // dt, dt=dt)

    freqs, A1, E1, _ = _gw2tdi(
        orbits_fn, gwfile, dt, tdi_generation=2, tdi_combination="AET"
    )

    ## compute fast approx
    fgb = JaxGB(orbits=_orbits, t_obs=duration, t0=t0, n=size)

    A2, E2, _ = fgb.get_tdi(params, tdi_generation=2, tdi_combination="AET")
    kmin = fgb.get_kmin(params[0, 0])  # frequency

    ## align on A/E
    df = 1 / fgb.t_obs
    i0 = np.where(freqs == df * kmin)[0][0]
    sl = slice(i0, i0 + A2.shape[1])
    olap = overlap(np.hstack([A1[sl], E1[sl]]), np.hstack([A2, E2]))
    assert olap > 0.99
