import torch


def final_displacement_error(pred_pos, pred_pos_gt, consider_ped=None, mode='sum'):
    """
    Computes the euclidean displacement error between two positions, assumed the be the final positions of trajectories.
    :param pred_pos: Tensor of shape (batch, 2). Predicted last position
    :param pred_pos_gt: Tensor of shape (batch, 2). Ground truth last position
    :param consider_ped: Tensor of shape (batch) -> which pedestrians to consider (1 to consider, 0 otherwise; or
    possibly decimal values if we want to give more contribution to some pedestrians)
    :param mode: Can be one of sum, raw
    :return: the Euclidean displacement error for this last position
    """
    loss = pred_pos_gt - pred_pos
    loss = loss ** 2
    if consider_ped is not None:
        loss = torch.sqrt(loss.sum(dim=1)) * consider_ped
    else:
        loss = torch.sqrt(loss.sum(dim=1))
    if mode == 'raw':
        return loss
    else:
        return torch.sum(loss)