import random


def get_train_val_idx(num_graphs, cv_fold):
    """ Return a tuple of the train and val indices,
    depending on the cv_fold
    This method shuffles the index (with a seed)
    The shuffle is consistent across machines with python3"""
    #
    # Extract indices of train and val in terms of the shuffled list
    # Balanced across test and train
    # Assumes 10-class
    #
    random.seed(1)
    num_classes = 10
    num_per_class = int(num_graphs/num_classes)
    val_size = int(0.2 * num_per_class)
    idx_to_classes = {}
    val_idx = []
    train_idx = []
    for cc in range(num_classes):
        idx_to_classes[cc] = list(range(cc*num_per_class, (cc+1)*num_per_class))
        random.shuffle(idx_to_classes[cc])
        # These indices correspond to the validation for this class.
        class_val_idx = slice(cv_fold * val_size, cv_fold * val_size + val_size, 1)
        # Extract validation.
        vals = idx_to_classes[cc][class_val_idx]
        val_idx.extend(vals)
        train_idx.extend(list(set(idx_to_classes[cc]) - set(vals)))
    #
    return tuple(train_idx), tuple(val_idx)