import torch


def collate_fn(data):
    """Pad the sequences.

    This function will be used to pad the sessions to max length
    in the batch and transpose the batch from
    batch_size x max_seq_len to max_seq_len x batch_size.
    It will return padded vectors, labels and lengths of each session (before padding)
    It will be used in the Dataloader.

    Args:
        data (pytorch Dataset): Sequential dataset.

    Returns:
        padded_sesss (Tensor): Padded vectors.
        labels (Tensor): Target item.
        lens (list): Lengths of each padded vector.
    """
    data.sort(key=lambda x: len(x[0]), reverse=True)
    lens = [len(sess) for sess, label in data]
    labels = []
    padded_sesss = torch.zeros(len(data), max(lens)).long()
    for i, (sess, label) in enumerate(data):
        padded_sesss[i, : lens[i]] = torch.LongTensor(sess)
        labels.append(label)
    padded_sesss = padded_sesss.transpose(0, 1)
    return padded_sesss, torch.tensor(labels).long(), lens