import torch


def mle_loss(uv_hat, uv, cov):
    """
    Calculate the MLE loss 
    :param uv_hat: reprojected uv shape(..., 2)
    :param uv:     UV measurements shape(..., 2)
    :param cov:    2x2 covariance matrix for uv shape(..., 2, 2)
    """
    res = uv - uv_hat

    # Ensure cov invertable
    cov[...,0,0] += 1e-6
    cov[...,1,1] += 1e-6

    # Ensure stability with cholesky decomp 
    # (NOTE can cause crash https://github.com/pytorch/pytorch/pull/50957)
    #S = torch.cholesky(torch.inverse(cov))
    #resS = torch.squeeze(torch.unsqueeze(res, -2) @ S, -2)
    #loss_mv = resS.square().sum(-1).mean()
    loss_mv = (res.unsqueeze(-2) @ torch.inverse(cov) @ res.unsqueeze(-1)).mean()

    # Add log determinant of cov to complete MLE loss
    # Note that the 0.5 multiplication happens in the main loss function
    def logdet_2x2(x):
        det = x[...,0,0]*x[...,1,1] - x[...,1,0]*x[...,0,1]
        ld = torch.log(torch.maximum(det, 1e-12*torch.ones_like(det)))
        
        # NOTE torch.logdet can fail because of a MAGMA issue.
        #ld_ = torch.logdet(x)
        #assert torch.allclose(ld, ld_)
        return ld

    loss_cov = logdet_2x2(cov).mean()

    return loss_mv, loss_cov