import numpy as np
import pytest

import scgenerator as sc


def test_pulse_envelope():
    t = np.linspace(-100, 100, 2049)

    assert np.trapz(sc.pulse.pulse_envelope(t, 10, energy=3) ** 2, x=t) == pytest.approx(3)
    assert (sc.pulse.pulse_envelope(t, 10, peak_power=3) ** 2).max() == pytest.approx(3)


def test_chirp():
    t = np.linspace(-10, 10, 1024)
    f = sc.pulse.gaussian_pulse(t, 1, 1, chirp=1.0)
    c = np.unwrap(np.angle(f))
    c -= c.min()
    assert pytest.approx(c, rel=1e-2, abs=1e-4) == t**2

    params = sc.Parameters(
        time_window=20,
        t_num=1024,
        peak_power=1,
        wavelength=800e-9,
        shape="gaussian",
        chirp=1,
        t0=1,
        quantum_noise=False,
    )
    t2, f2 = params.compute("t", "field_0")
    c2 = np.unwrap(np.angle(f2))
    c2 -= c2.min()
    assert pytest.approx(c2, rel=1e-2, abs=1e-4) == t2**2
    assert pytest.approx(sc.abs2(f2)) == 16 ** (-((t2 / np.sqrt(2 * np.log(2))) ** 2))