import numpy as np
from giant.data.dataset import make_event_split


def test_make_event_split_sizes():
    rng = np.random.default_rng(42)
    event_ids = rng.integers(0, 50, size=1000)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    unique = np.unique(event_ids)
    assert len(train_set) + len(val_set) == len(unique)


def test_make_event_split_no_overlap():
    rng = np.random.default_rng(7)
    event_ids = rng.integers(0, 50, size=1000)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    assert train_set.isdisjoint(val_set)


def test_make_event_split_no_empty_sets():
    rng = np.random.default_rng(0)
    event_ids = rng.integers(0, 20, size=500)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    assert len(train_set) > 0
    assert len(val_set) > 0


def test_make_event_split_reproducible():
    event_ids = np.arange(100)
    a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
    b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
    assert a_tr == b_tr
    assert a_val == b_val