import numpy as np
import pytest

from scgenerator.math import abs2, wspace
from scgenerator.physics.units import m_rads
from scgenerator.physics.pulse import find_lobe_limits
from scgenerator.spectra import Spectrum


def test_export():
    t = np.linspace(-1e-12, 1e-12, 1024)
    w = wspace(t) + m_rads(800e-9)

    spec = np.fft.fft(np.exp(-((t / 1e-13) ** 2)))
    spec = Spectrum(spec, w, t)

    new_spec = Spectrum.from_bytes(bytes(spec))
    assert np.all(new_spec.w == spec.w)
    assert np.all(new_spec.t == spec.t)
    assert np.all(new_spec.l == spec.l)
    assert np.all(new_spec == spec)
    assert new_spec.ifft is spec.ifft

    t = np.linspace(-1e-12, 1e-12, 512)
    w = np.fft.rfftfreq(len(t), d=t[1] - t[0]) * 2 * np.pi + m_rads(800e-9)
    spec = np.exp(2j * np.pi * np.random.rand(4, 2, 5, 257))
    spec = Spectrum(spec, w, t, np.fft.irfft)

    new_spec = Spectrum.from_bytes(bytes(spec))
    assert np.all(new_spec.w == spec.w)
    assert np.all(new_spec.t == spec.t)
    assert np.all(new_spec.l == spec.l)
    assert np.all(new_spec == spec)
    assert new_spec.ifft is spec.ifft


def test_center_gravity():
    t = np.linspace(-1e-12, 1e-12, 1024)
    w = wspace(t) + m_rads(800e-9)

    spec = np.fft.fft(np.exp(-((t / 1e-13) ** 2)))
    spec = Spectrum(spec, w, t)
    assert spec.center_of_gravity == pytest.approx(0)

    spec = np.fft.fft(
        np.array([np.exp(-((t / 1e-13) ** 2)), np.exp(-(((t - 1e-15) / 1e-13) ** 2))])
    )
    spec = Spectrum(spec, w, t)
    assert spec.center_of_gravity * 1e15 == pytest.approx((0, 1))

    spec = np.fft.fft(
        np.array(
            [
                [
                    np.exp(-((t / 1e-13) ** 2)),
                    np.exp(-(((t - 1e-15) / 1e-13) ** 2)),
                    np.exp(-(((t + 2e-15) / 1e-13) ** 2)),
                ],
                [
                    np.exp(-((t / 1e-13) ** 2)),
                    np.exp(-(((t + 1e-15) / 1e-13) ** 2)),
                    np.exp(-(((t - 2e-15) / 1e-13) ** 2)),
                ],
            ]
        )
    )
    spec = Spectrum(spec, w, t)
    assert spec.center_of_gravity[0] * 1e15 == pytest.approx((0, 1, -2))
    assert spec.center_of_gravity[1] * 1e15 == pytest.approx((0, -1, 2))


def test_energy():
    t = np.linspace(-1e-12, 1e-12, 1024)
    w = wspace(t) + m_rads(800e-9)

    spec = np.random.rand(1024)
    spec = Spectrum(spec, w, t)
    assert spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t)

    spec = np.random.rand(2, 1024)
    spec = Spectrum(spec, w, t)
    assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
    assert spec.energy.shape == (2,)

    spec = np.random.rand(3, 2, 1024)
    spec = Spectrum(spec, w, t)
    assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
    assert spec.energy.shape == (3, 2)


def test_mask():
    wl = 800e-9
    bw = 50e-9
    n = 8192
    t = np.linspace(-1e-12, 1e-12, n)
    w = wspace(t) + m_rads(wl)
    spec = Spectrum(np.ones(n) + 0j, w, t)
    mased = spec.mask_wl(wl, bw)
    l, r = find_lobe_limits(mased.wl_disp, mased.wl_int)[1]
    assert r - l == pytest.approx(bw, rel=1e-2, abs=1e-10)