import torch


def collate_train(batch):
    """
    Collate function for preparing the training batch for DCL training.

    :param batch: The list containing outputs of the __get_item__() function.
    Length of the list is equal to the required batch size.
    :return: The batch containing images and labels for DCL training
    """
    imgs = []  # List to store the images
    target = []  # List to store the targets
    target_jigsaw = []  # List to store the jigsaw targets
    patch_labels = []  # List to store the ground truth patch labels
    # Iterate over each output of the __get_item__() function
    for sample in batch:
        imgs.append(sample[0])  # Append the original image
        imgs.append(sample[1])  # Append the transformed image (after applying RCM)
        # Append the same labels for original and RCM image for Cls. head
        target.append(sample[2])
        target.append(sample[2])
        # Append the different labels for original and RCM image for Adv. head
        target_jigsaw.append(sample[2])
        target_jigsaw.append(sample[3])
        # Append the ground truth patch labels
        patch_labels.append(sample[4])
        patch_labels.append(sample[5])
    # Stack the images and return the required batch for training
    return torch.stack(imgs, 0), target, target_jigsaw, patch_labels