import torch


def get_pose_loss(gt_pose, pred_pose, vis, loss_type):
    """
    gt_pose: [B, J*2]
    pred_pose: [B, J*2]
    vis: [B, J]

    Loss: L2 (MSE), or L1
    """
    vis = torch.repeat_interleave(vis, 2, dim=-1)  # [B, 1122...JJ]

    if loss_type == "L2":
        loss = torch.sum((gt_pose - pred_pose) ** 2 * vis, dim=-1) / torch.sum(vis, dim=-1)  # [B]
    elif loss_type == "L1":
        loss = torch.sum(torch.abs(gt_pose - pred_pose) * vis, dim=-1) / torch.sum(vis, dim=-1)  # [B]
    else:
        assert False, "Unknown loss type"

    return loss