import numpy as np
from tidyfit.cv import kfold_indices, stratified_kfold_indices


def test_kfold_indices():
    folds = list(kfold_indices(10, n_splits=5, random_state=1))
    assert len(folds) == 5
    train, val = folds[0]
    assert len(set(train).intersection(set(val))) == 0


def test_stratified_kfold_indices():
    y = np.array([0, 0, 0, 1, 1, 1, 1, 0, 1, 0])
    folds = list(stratified_kfold_indices(y, n_splits=5, random_state=1))
    assert len(folds) == 5
    for train, val in folds:
        assert len(set(train).intersection(set(val))) == 0
