import numpy as np

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])
    )