"""Unit tests for module shift_inversion_dask"""

import numpy as np
import pytest

from lisainstrument.dynamic_delay_dask import numpyfy_dask_multi
from lisainstrument.shift_inversion_dask import make_shift_inverse_lagrange_dask
from lisainstrument.shift_inversion_numpy import make_shift_inverse_lagrange_numpy


def test_shift_inversion_dask():
    """Test basic functioning of shift_inversion_dask on analytic coordinate transform"""
    chunks = 70
    order = 31
    nsamp = 3000
    fsample = 16.0
    dt = 1 / fsample

    f_mod = 0.005
    a_mod = 1e-2 / (2 * np.pi * f_mod)
    max_it = 5
    tol = 1e-10

    def dx_from_x(x):
        return np.sin(2 * np.pi * f_mod * x) * a_mod

    xi = np.arange(nsamp) * dt

    dxi = dx_from_x(xi)

    op_np = make_shift_inverse_lagrange_numpy(
        order=order,
        fsample=fsample,
        max_abs_shift=a_mod * 1.01,
        max_iter=max_it,
        tolerance=tol,
    )
    op_da = make_shift_inverse_lagrange_dask(
        order=order,
        fsample=fsample,
        max_abs_shift=a_mod * 1.01,
        max_iter=max_it,
        tolerance=tol,
    )
    op_na = numpyfy_dask_multi(op_da, chunks)

    ai_np = op_np(dxi)
    ai_na = op_na(dxi)

    assert ai_np == pytest.approx(ai_na, abs=tol, rel=0)