import copy


def kfold_cross_validation(X, n_splits=5):
    """Split dataset into cross validation folds.

    Args:
        X(list of list of obj): The list of samples
            The shape of X is (n_samples, n_features)
        n_splits(int): Number of folds.

    Returns:
        X_train_folds(list of list of int): The list of training set indices for each fold
        X_test_folds(list of list of int): The list of testing set indices for each fold

    Notes: 
        The first n_samples % n_splits folds have size n_samples // n_splits + 1, 
            other folds have size n_samples // n_splits, where n_samples is the number of samples
            (e.g. 11 samples and 4 splits, the sizes of the 4 folds are 3, 3, 3, 2 samples)
        Loosely based on sklearn's KFold split(): https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.KFold.html
    """
    X_train_folds = []
    X_test_folds = []
    
    # We're going to make a list of the split-up folds to start.
    n = 0
    folds = []
    for i in range(n_splits):
        folds.append([])
    
    # Just get the indices
    while n < len(X):
        folds[n % n_splits].append(n)
        n += 1
        
    # Now we need to use these folds to generate n_splits sets for each
    for i in range(n_splits):
        new_train_fold = []
        new_test_fold = []
        for j in range(n_splits):
            if i == j:
                for instance in folds[j]:
                    new_test_fold.append(copy.deepcopy(instance))
            else:
                for instance in folds[j]:
                    new_train_fold.append(copy.deepcopy(instance))
                    
        X_train_folds.append(copy.deepcopy(new_train_fold))
        X_test_folds.append(copy.deepcopy(new_test_fold))

    return X_train_folds, X_test_folds