import torch


def collate_fn(batch):
    """
    Create a batch of data given a list of N sequences and labels. Sequences are stacked into a single tensor
    of shape (N, max_sequence_length), where max_sequence_length is the maximum length of any sequence in the
    batch. Sequences shorter than this length should be filled up with 0's. Also returns a tensor of shape (N, 1)
    containing the label of each sequence.

    :param batch: A list of size N, where each element is a tuple containing a sequence tensor and a single item
    tensor containing the true label of the sequence.

    :return: A tuple containing two tensors. The first tensor has shape (N, max_sequence_length) and contains all
    sequences. Sequences shorter than max_sequence_length are padded with 0s at the end. The second tensor
    has shape (N, 1) and contains all labels.
    """
    # print(batch)
    sentences, labels = zip(*batch)
    sequence_lenes = [len(x) for x in sentences]
    max_sequence_length = max(sequence_lenes)
    # print("max_sequence_length: ", max_sequence_length)

    data_matrix = torch.zeros((len(sentences), max_sequence_length))
    for i in range(len(sentences)):
        for j in range(sequence_lenes[i]):
            data_matrix[i, j] = sentences[i][j]
    # print("data_matrix shape: ", data_matrix.shape)

    return (data_matrix, torch.tensor(labels))