import torch


def siamese_collate(batch):
    """
    This collator is used in Jigsaw approach.
    Input:
        batch: Example
                batch = [
                    {"data": [img1,], "label": [lbl1, ]},        #img1
                    {"data": [img2,], "label": [lbl2, ]},        #img2
                    .
                    .
                    {"data": [imgN,], "label": [lblN, ]},        #imgN
                ]
                where:
                    img{x} is a tensor of size: num_towers x C x H x W
                    lbl{x} is an integer
    Returns: Example output:
                torch.tensor([img1_0, ..., imgN_0]),
                torch.tensor([lbl1, ..., lblN])

                where the output is of dimension: (N * num_towers) x C x H x W

    see https://github.com/facebookresearch/vissl/blob/aa3f7cc33b3b7806e15593083aedc383d85e4a53/vissl/data/collators/siamese_collator.py
    """
    assert "data" in batch[0], "data not found in sample"
    assert "label" in batch[0], "label not found in sample"
    num_data_sources = 1
    batch_size = len(batch)
    data = [x["data"] for x in batch]
    labels = [x["label"] for x in batch]

    # each image is of shape: num_towers x C x H x W
    # num_towers x C x H x W -> N x num_towers x C x H x W
    idx_data = torch.stack([data[i] for i in range(batch_size)])
    idx_labels = [labels[i] for i in range(batch_size)]
    batch_size, num_siamese_towers, channels, height, width = idx_data.size()
    # N x num_towers x C x H x W -> (N * num_towers) x C x H x W
    idx_data = idx_data.view(batch_size * num_siamese_towers, channels, height, width)
    should_flatten = False
    if idx_labels[0].ndim == 1:
        should_flatten = True
    idx_labels = torch.stack(idx_labels).squeeze()
    if should_flatten:
        idx_labels = idx_labels.flatten()

    return idx_data, idx_labels