import numpy as np
import pytest

import scgenerator as sc


def test_normalisation():
    rng = np.random.default_rng(56)
    t = np.linspace(-10, 10, 512)
    s = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15

    target = np.sum(sc.abs2(np.fft.fft(s))) / 512
    noise = sc.noise.NoiseMeasurement.from_time_series(s, 1, "boxcar", detrend=False)
    assert np.sum(noise.psd) == pytest.approx(target)


def test_time_and_back():
    """
    sampling a time series from a spectrum and transforming
    it back yields the same spectrum
    """
    rng = np.random.default_rng(57)
    t = np.linspace(-10, 10, 512)
    signal = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15

    noise = sc.noise.NoiseMeasurement.from_time_series(signal, 1, "boxcar", detrend=False)
    _, new_signal = noise.time_series(len(noise.freq))
    new_noise = sc.noise.NoiseMeasurement.from_time_series(new_signal, 1, "boxcar", detrend=False)
    assert new_noise.psd == pytest.approx(noise.psd)


def test_nyquist():
    """
    generating a time series and tranforming it back yields the same spectrum.
    Using segements, the nyquist frequency is exactly spread out over the frequency bin width
    """
    signal = np.cos(np.arange(1024) * np.pi)

    n1 = sc.noise.NoiseMeasurement.from_time_series(signal, 1, None)
    n3 = sc.noise.NoiseMeasurement.from_time_series(signal, 1, None, 512)
    n15 = sc.noise.NoiseMeasurement.from_time_series(signal, 1, None, 128)

    assert n1.psd[-1] == n3.psd[-1] * 2 == n15.psd[-1] * 8


def test_sampling():
    f = np.geomspace(10, 2e6, 138)
    spec = 1 / (f + 1)

    noise = sc.noise.NoiseMeasurement(f, spec)

    assert noise.sample_spectrum(257)[0][0] == 0
    assert noise.sample_spectrum(257, log_mode=True)[0][0] == 0