import torch


def _collate_fn(batch):
  """
  Custom collate function to manually specify how samples are batched from the DataLoader.
  See https://pytorch.org/tutorials/beginner/data_loading_tutorial.html#iterating-through-the-dataset
  for an introduction.

  In particular, the batch represents a collection of
    (lmk_seq, caption)
  pairs. However each sequence may be of different sequence length.

  :param batch: The complete batch of data samples to be collated. In particular, this is a list of frames and captions,
    each of shape (seq_len, num_pts, pt_dim), (num_chars)
  """
  def _pad(seqs, dtype=torch.float32):
    """ Pads a batch of sequences of varying seq_len. """
    assert len(seqs) > 0 and all(x.shape[1:] == seqs[0].shape[1:] for x in seqs)
    lens = torch.LongTensor([len(x) for x in seqs])
    max_seq_len = torch.max(lens)

    # padded_seq_dims: (batch, max_seq_len, ...).
    padded_seq_dims = (len(seqs), max_seq_len,) + seqs[0].shape[1:]
    res = torch.zeros(padded_seq_dims, dtype=dtype)
    for i, seq in enumerate(seqs):
      src_len = lens[i]
      res[i, :src_len] = torch.Tensor(seq)
    return res, lens

  assert all(len(x) == 2 for x in batch)
  # (1, batch, (seq_len, 68, 3))
  frames, captions = zip(*batch)

  # Merge sequences (from tuple of 1D tensor to 2D tensor)
  # (batch, seq_len, ...)
  src_seqs, src_lens = _pad(frames, dtype=torch.float32)
  tgt_seqs, tgt_lens = _pad(captions, dtype=torch.long)
  return src_seqs, src_lens, tgt_seqs, tgt_lens