import torch


def kp_l2_loss(kp_pred, kp_gt, reduction='mean'):
    """
    L2 loss between visible keypoints.

    \Sum_i [0.5 * vis[i] * (kp_gt[i] - kp_pred[i])^2] / (|vis|)
    """
    # criterion = torch.nn.MSELoss()
    criterion = torch.nn.L1Loss(reduction='none')

    vis = (kp_gt[:, :, 2] > 0).float()
    loss = criterion(kp_pred, kp_gt[:, :, :2]).sum(-1) * vis
    loss = loss.mean(-1) / (vis.mean(-1) + 1e-4)
    if reduction == 'mean':
        return loss.mean()
    else:
        return loss