import pytest

import scgenerator as sc

PARAMS1 = sc.Parameters(
    name="Dudley2006",
    wavelength=835e-9,
    width=50e-15,
    peak_power=10e3,
    repetition_rate=20e6,
    shape="sech",
    # fiber
    length=15e-2,
    raman_type="measured",
    beta2_coefficients=[
        -11.830e-27,
        8.1038e-42,
        -9.5205e-57,
        2.0737e-71,
        -5.3943e-86,
        1.3486e-100,
        -2.5495e-115,
        3.0524e-130,
        -1.7140e-145,
    ],
    gamma=0.11,
    # simulation
    tolerated_error=1e-6,
    wavelength_window=[400e-9, 1500e-9],
    t_num=8192,
    quantum_noise=False,
    z_num=128,
)
PARAMS2 = sc.Parameters(
    name="dual comb",
    wavelength=1056e-9,
    width=78e-15,
    peak_power=13e3,
    repetition_rate=20e6,
    shape="gaussian",
    # fiber
    length=20e-2,
    n2=2.2e-20,
    raman_type="measured",
    dispersion_file="tests/data/PM1050NEG_dispersion.npz",
    loss_file="tests/data/PM1050NEG_loss.npz",
    effective_area_file="tests/data/PM1050NEG_aeff_slow.npz",
    # simulation
    interpolation_degree=12,
    tolerated_error=1e-8,
    wavelength_window=[500e-9, 1500e-9],
    t_num=8192,
    quantum_noise=True,
    z_num=128,
)


def test_poly_dispersion_grid_size():
    params = PARAMS1.compile()
    disp_ind = params.compute("dispersion_ind")

    assert disp_ind is not None
    assert all((params.l[disp_ind] < 1500e-9) & (params.l[disp_ind] > 400e-9))
    assert len(params.compute("beta2_arr")) == params.t_num


def test_custom_dispersion_grid_size():
    params = PARAMS2.compile()
    disp_ind = params.compute("dispersion_ind")

    assert disp_ind is not None
    assert all((params.l[disp_ind] < 1500e-9) & (params.l[disp_ind] > 400e-9))
    assert len(params.compute("beta2_arr")) == params.t_num