import torch


def l2_loss(pred_traj, pred_traj_gt, mode='average', type = "mse"):
    """
    Input:
    - pred_traj: Tensor of shape (seq_len, batch, 2). Predicted trajectory.
    - pred_traj_gt: Tensor of shape (seq_len, batch, 2). Groud truth
    predictions.
    - mode: Can be one of sum, average, raw
    Output:
    - loss: l2 loss depending on mode
    """
    # disabled loss mask
    #loss_mask = 1
    seq_len, batch, _ = pred_traj.size()
    #loss_mask = loss_mask.permute(2, 1,0)[:, :, 0]
    d_Traj = pred_traj_gt - pred_traj

    if type =="mse": 
        loss = torch.norm( (d_Traj), 2, -1)
    elif type =="average":
      
        loss = (( torch.norm( d_Traj, 2, -1)) + (torch.norm( d_Traj[-1], 2, -1)))/2.
       
    else: 
        raise AssertionError('Mode {} must be either mse or  average.'.format(type))


    if mode == 'sum':
        return torch.sum(loss)
    elif mode == 'average':
     
        return torch.mean(loss, dim =0 )
    elif mode == 'raw':
        return loss.sum(dim=0)