import random
import json


def generate_train_val_indices(input_size, val_size, file_name):
    """Generate random indices for training and validation set

    Parameters
    ==========
    input_size : int
        Dataset size.
    val_size : int
        Validation size.
    file_name : str
        Output json file name.

    Returns
    =======
    train_indices : list(int)
    val_indices : list(int)
    """
    assert (val_size < input_size)
    indices = [x for x in range(0, input_size)]
    random.shuffle(indices)
    train_indices = indices[val_size:]
    validation_indices = indices[:val_size]
    print("train size {} val size {}".format(
        len(train_indices), len(validation_indices)))

    # Write indices to a json file
    with open(file_name, 'w') as fp:
        data = {'train': train_indices, 'validation': validation_indices}
        json.dump(data, fp)
        print("Saved to {}...".format(file_name))

    return train_indices, validation_indices