import torch


def rotMat_to_quat(rotMat):
    """Convert rotation matrices back to quaternions.

    Ported from Matlab:
    https://github.com/asheshjain399/RNNexp/blob/srnn/structural_rnn/CRFProblems/H3.6m/mhmublv/Motion/rotmat2quat.m#L4


    Args:
        rotMat (torch.Tensor): rotation matrix or matrices to convert to
            quaternion format.

    Returns:
        torch.Tensor: quaternion(s) converted from rotation matrix format
    """
    if len(rotMat.shape) != 3:
        rotMat = rotMat.unsqueeze(0)

    assert rotMat.shape[1] == 3 and rotMat.shape[2] == 3

    diffMat = rotMat - torch.transpose(rotMat, 1, 2)

    r = torch.zeros((rotMat.shape[0], 3), dtype=torch.float64)

    r[:, 0] = -diffMat[:, 1, 2]
    r[:, 1] = diffMat[:, 0, 2]
    r[:, 2] = -diffMat[:, 0, 1]

    sin_theta = torch.norm(r, dim=1)/2
    sin_theta = sin_theta.unsqueeze(1)

    r0 = r / (torch.norm(r, dim=1).unsqueeze(1) + 1e-9)

    cos_theta = (rotMat.diagonal(dim1=-2, dim2=-1).sum(-1) - 1) / 2
    cos_theta = cos_theta.unsqueeze(1)

    theta = torch.atan2(sin_theta, cos_theta)

    theta = theta.squeeze(1)

    q = torch.zeros((rotMat.shape[0], 4), dtype=torch.float64)

    q[:, 0] = torch.cos(theta/2)
    q[:, 1:] = r0*torch.sin(theta/2).unsqueeze(1)

    return q