import numpy as np
import pytest

import scgenerator as sc


def test_segmentation():
    t = np.arange(32)

    r = np.arange(16)
    assert np.all(sc.noise.segments(t, 3) == np.vstack([r, r + 8, r + 16]))

    r = np.arange(8)

    assert np.all(sc.noise.segments(t, 4) == np.vstack([r, r + 6, r + 12, r + 18]))
    assert np.all(sc.noise.segments(t, 5) == np.vstack([r, r + 5, r + 10, r + 15, r + 20]))
    assert np.all(sc.noise.segments(t, 6) == np.vstack([r, r + 4, r + 8, r + 12, r + 16, r + 20]))
    assert np.all(
        sc.noise.segments(t, 7) == np.vstack([r, r + 4, r + 8, r + 12, r + 16, r + 20, r + 24])
    )


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, "Square", force_no_dc=False)
    assert np.sum(noise.psd) == pytest.approx(target)


def test_no_dc():
    rng = np.random.default_rng(56)
    t = np.linspace(-10, 10, 512)
    s = rng.normal(0, 1, len(t)).cumsum()
    noise = sc.noise.NoiseMeasurement.from_time_series(s, 1)
    assert noise.psd[0] == pytest.approx(0)


def test_time_and_back():
    """
    sampling a time series from a spectrum and transforming
    it back should yield 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, "Square", force_no_dc=False)
    _, new_signal = noise.time_series(len(t))
    new_noise = sc.noise.NoiseMeasurement.from_time_series(
        new_signal, 1, "Square", force_no_dc=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, 1)
    n3 = sc.noise.NoiseMeasurement.from_time_series(signal, 1, None, 3)
    n15 = sc.noise.NoiseMeasurement.from_time_series(signal, 1, None, 15)

    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