from sklearn.model_selection import KFold, ShuffleSplit


def cross_val(model, data, hyperparams=None, n_splits=None, test_size=None, cbs=None):
    arg_err = """Either n_splits for cross validation or
    test_size for shuffle split must be specified."""
    assert (n_splits, test_size) != (None, None), arg_err
    if n_splits:
        kfold = KFold(n_splits, random_state=0)
    else:
        kfold = ShuffleSplit(n_splits, test_size=test_size, random_state=0)
    losses, histories = [], []
    features, labels = data

    for i, (train_index, test_index) in enumerate(kfold.split(features)):
        model.log(f"cross validation fold {i + 1}/{n_splits}")
        train = (features[train_index], labels[train_index])
        test = (features[test_index], labels[test_index])
        loss = model.run((train, test), hyperparams=hyperparams, cbs=cbs)
        losses.append(loss)
        histories.append(model.history.history["loss"])

    return sum(losses) / len(losses)
