import torch


def extract_patches(arguments, images, num_patches):
    """
    Extracts a number of patches from a given tensor of images.
    :param arguments: Dictionary containing 'cpc_patch_size' and 'cpc_patch_stride'.
    :param images: A tensor of input images.
    :param num_patches: Integer for the number of patches to take from each dimension.
    :return: A tensor of patches equal to 'num_patches * num_patches'.
    """

    # List of all patches.
    all_patches = []

    # Cycles through all the inputted images.
    for i in range(images.shape[0]):

        # Cycles through the patches in the image.
        for y_patch in range(num_patches):
            for x_patch in range(num_patches):

                # Gets the y coordinates for the patch.
                y1 = y_patch * arguments["cpc_patch_stride"]
                y2 = y1 + arguments["cpc_patch_size"]

                # Gets the x coordinates for the patch.
                x1 = x_patch * arguments["cpc_patch_stride"]
                x2 = x1 + arguments["cpc_patch_size"]

                # Extracts the patch from the image.
                img_patch = images[i, :, y1:y2, x1:x2]
                img_patch = img_patch.view(images.shape[1], arguments["cpc_patch_size"], arguments["cpc_patch_size"])
                all_patches.append(img_patch)

    # Concatenates all the patches into a single tensor of images.
    all_patches = torch.cat(all_patches)
    return all_patches.view([-1, 3, arguments["cpc_patch_size"], arguments["cpc_patch_size"]])