import random


def make_accuracy_data(num_classes, num_test_samples, num_correct_predictions, k):
    """Convenience method to create the labels and predictions array.
    Args:
        num_classes (int): The number of classes in the dataset.
        num_test_samples (int): The number of labels and predictions to generate.
        num_correct_predictions (int): The number of predictions that are correct.
        k (int): The number of class labels in a prediction.
    Returns:
        list : The labels.
        list : The predictions.
    """

    assert k <= num_classes, "k too big"
    assert num_correct_predictions <= num_test_samples, ""\
        "`num_correct_predictions` is larger than `num_test_samples`"

    if k == num_classes:
        assert num_test_samples == num_correct_predictions, ""\
            "`num_correct_predictions` should be equal to `num_test_samples` "\
            "when `k` equals `num_classes`"

    labels = []
    predictions = []

    class_labels = range(num_classes)

    # Determine which idexes will be correct
    if num_correct_predictions > 0:
        correct_prediction_idxs = set(random.sample(range(num_test_samples),
                                                    num_correct_predictions))
    else:
        correct_prediction_idxs = set()

    # Fill in the labels and prediction lists
    for i in range(num_test_samples):
        gt_label = random.choice(class_labels)
        labels.append(gt_label)

        preds = random.sample(class_labels, k)
        if i in correct_prediction_idxs:
            if gt_label not in preds:
                preds[0] = gt_label
        else:
            if gt_label in preds:
                preds = [p for p in preds if p != gt_label]
                preds.append(gt_label + 1 % num_classes)
        predictions.append(preds)

    return labels, predictions