import torch


def prepare_batch(batch):
    """This is the collate function.

    The mass spectra must be padded so that they fit nicely as a tensor.
    However, the padded elements are ignored during the subsequent steps.

    Parameters
    ----------
    batch : tuple of tuple of torch.tensor
        A batch of data from a PairedSpectrumDataset.

    Returns
    -------
    X : torch.Tensor of shape (batch_size, n_peaks, 2)
    Y : torch.Tensor of shape (batch_size, n_peaks, 2)
        The mass spectra to embed, where ``X[:, :, 0]`` are the m/z values
        and ``X[:, :, 1]`` are their associated intensities.
    gsp : torch.Tensor of shape (batch_size,)
        The GSP calculated between the mass spectra in X and Y.
    """
    X, Y, gsp = list(zip(*batch))
    X = torch.nn.utils.rnn.pad_sequence(X, batch_first=True)
    Y = torch.nn.utils.rnn.pad_sequence(Y, batch_first=True)
    gsp = torch.tensor(gsp)
    return X, Y, gsp