import torch


def seq_collate(data):
    """
    Collect the several fields of the (collate <-> collect and combine)
    :param data: the original data, not formatted
    :return: a tuple containing the several parts of the data:
        - obs_traj: Tensor of shape (obs_seq_len, batch, 2). The observed trajectory, in absolute coordinates
        - pred_traj: Tensor of shape (pred_seq_len, batch, 2). The predicted trajectory, in absolute coordinates
        - obs_traj_rel: Tensor of shape (obs_seq_len, batch, 2). The observed trajectory, in relative displacements
        - pred_traj_rel: Tensor of shape (pred_seq_len, batch, 2). The predicted trajectory, in relative displacements
        - non_linear_ped: Tensor of shape (batch). 1 if trajectory is non linear, 0 otherwise
        - loss_mask: Tensor of shape (traj_seq_len, batch). Loss mask to (optionally) use in training.
        - seq_start_end: Tensor containing the beginning and end (in terms of pedestrians in this batch) of a sequence.
            A sequence corresponds to a set of trajectories that belong to the same time frame. Note that the batch may
            have several sequences.
    """
    (obs_seq_list, pred_seq_list, obs_seq_rel_list, pred_seq_rel_list, metadata_list, loss_mask_list,
     seq_start_end) = zip(*data)

    # Data format: batch, input_size, seq_len
    # LSTM input format: seq_len, batch, input_size
    obs_traj = torch.cat(obs_seq_list, dim=0)
    pred_traj = torch.cat(pred_seq_list, dim=0)
    obs_traj_rel = torch.cat(obs_seq_rel_list, dim=0)
    pred_traj_rel = torch.cat(pred_seq_rel_list, dim=0)
    loss_mask = torch.cat(loss_mask_list, dim=0)
    seq_start_end = torch.cat(seq_start_end, dim=0)
    out = [obs_traj, pred_traj, obs_traj_rel, pred_traj_rel, metadata_list, loss_mask, seq_start_end]

    return tuple(out)