import torch


def l2_loss(pred_fut_traj, fut_traj, mode="average"):
    """
    Compute L2 loss

    Input:
    - pred_fut_traj: Tensor of shape (seq_len, batch, 2). Predicted trajectory.
    - fut_traj: Tensor of shape (seq_len, batch, 2). Groud truth future trajectory.
    - mode: Can be one of sum, average, raw
    Output:
    - loss: l2 loss depending on mode
    """

    loss = (fut_traj[:,:,:2].permute(1, 0, 2) - pred_fut_traj.permute(1, 0, 2)) ** 2
    if mode == "sum":
        return torch.sum(loss)
    elif mode == "average":
        return torch.mean(loss) 
    elif mode == "raw":
        return loss.sum(dim=2).sum(dim=1)