import torch


def l2_loss(pred_traj, pred_traj_gt, loss_mask, random=0, mode='average'):
    """
    Computes standard L2 norm loss between predicted and ground truth trajectories. This involves squaring the
    differences between coordinates, and summing over all length (pred_traj_len)
    :param pred_traj: Tensor of shape (pred_traj_len, batch, 2). Predicted trajectory.
    :param pred_traj_gt: Tensor of shape (pred_traj_len, batch, 2). Ground truth trajectory to compare with prediction
    :param loss_mask: Tensor of shape (batch, traj_len). Applies a mask to the loss values (in case one doesn't want to
    consider some trajectories for the sake of loss computation).
    :param random: not used
    :param mode: Can be one of sum, average, raw
    :return: l2 loss depending on mode (tensor can have different dimensions)
    """
    traj_len, batch, _ = pred_traj.size()
    # switch to shape (batch, traj_len, 2)
    loss = (pred_traj_gt.permute(1, 0, 2) - pred_traj.permute(1, 0, 2)) ** 2
    if mode == 'sum':
        return torch.sum(loss)
    elif mode == 'average':
        return torch.sum(loss) / torch.numel(loss[:, :, 0].data)
    elif mode == 'raw':
        return loss.sum(dim=2).sum(dim=1)
    else:
        # different kind of mode, returning per each instant
        return loss.sum(dim=2).permute(1, 0)