import torch


def collate_fn(data):
    """ Creates mini-batch from x, ivec, jvec tensors

    We should build custom collate_fn, as the ivec, and jvec have varying lengths. These should be appended
    in row form

    Args:
        data: list of tuples contianing (x, ivec, jvec)

    Returns:
        x: one hot encoded vectors stacked vertically
        ivec: long vector
        jvec: long vector
    """
    x, m, ivec, jvec, demo = zip(*data)
    m = torch.stack(m, dim=1)
    x = torch.stack(x, dim=1)
    ivec = torch.cat(ivec, dim=0)
    jvec = torch.cat(jvec, dim=0)
    demo = torch.stack(demo, dim=1)
    return x, m, ivec, jvec, demo