import torch


def get_kth_fold_data(X_train, y_train, k, iteration):
    """ Function used to extract the datasets for the k-th iteration of the k-fold cross validation algorithm """
    # determine how many examples each fold gets
    portion = X_train.shape[0] // k

    # distinguish the cases and act accordingly
    if iteration == 0:
        X_val = X_train[0 : portion]
        y_val = y_train[0 : portion]

        _X_train = X_train[portion:]
        _y_train = y_train[portion:]

    elif iteration < k - 1:
        X_val = X_train[iteration * portion : (iteration + 1) * portion]
        y_val = y_train[iteration * portion : (iteration + 1) * portion]

        left_X = X_train[:iteration * portion]
        right_X = X_train[(iteration + 1) * portion:]
        _X_train = torch.cat((left_X, right_X), 0)

        left_y = y_train[:iteration * portion]
        right_y = y_train[(iteration + 1) * portion:]
        _y_train = torch.cat((left_y, right_y), 0)

    else:
        X_val = X_train[-portion:]
        y_val = y_train[-portion:]

        _X_train = X_train[:-portion]
        _y_train = y_train[:-portion]


    # return the target datasets
    return _X_train, _y_train, X_val, y_val