import torch


def collate_fn(batch):
    """Creates mini-batch tensors from the list of items
        (image, question_label, question_length, answer).

    We should build custom collate_fn rather than using default collate_fn,
    because merging question (including padding) is not supported in default.
    Args:
        batch: list of two components [network inputs, image info]
            - network intpus: list of inputs [imgs, qst_labels, qst_lengths,
            (assignments), answers]
            - img_info: question ids or image filenames
    Returns: Sorted items in decreasing order of question length
        imgs: torch FloatTensor [batch_size, feat_dim, kernel_size, kernel_size].
        qst_labels: torch LongTensor [batch_size, padded_length].
        qst_lengths: list whose length is batch_size; valid length for each padded caption.
        (assignments): pre-computed assignments
        answers: torch LongTensor [batch_size]
    """
    # sort a batch by question length (descending order)
    batch.sort(key=lambda x: x[2], reverse=True)
    batch = [val for val in zip(*batch)]

    items = batch[:-1]
    img_info = batch[-1] # qst_id or img_filename
    return_vals = [items, img_info]

    # stacking tensors
    num_items = len(items)
    items[0] = torch.stack(items[0], 0) # images
    items[1] = torch.stack(items[1], 0).long() # question labels
    items[-1] = torch.stack(items[-1], 0).squeeze() # answers
    if num_items == 5:
        items[-2] = torch.stack(items[-2], 0) # for pre-computed logits of base models
    return return_vals