import torch


def rotmat_to_quat(rotmat):
    """Convert rotation matrix to quaternion coefficients.
    Args:
        rotmat: size is [B, 3, 3]
    Returns:
        Quaternion: size is [B, 4] <===> (w, x, y, z)
    """
    quaternion = torch.zeros([rotmat.size(0), 4], device=rotmat.device)
    trace = rotmat[:, 0, 0] + rotmat[:, 1, 1] + rotmat[:, 2, 2]
    flag = 1 + trace > 0
    s = torch.zeros_like(trace)

    # pos
    s[flag] = 2 * torch.sqrt(1 + trace[flag]) + 1e-16
    s_pos = s[flag]
    quaternion[flag, 0] = s_pos / 4
    quaternion[flag, 1] = (rotmat[flag, 2, 1] - rotmat[flag, 1, 2]) / s_pos
    quaternion[flag, 2] = (rotmat[flag, 0, 2] - rotmat[flag, 2, 0]) / s_pos
    quaternion[flag, 3] = (rotmat[flag, 1, 0] - rotmat[flag, 0, 1]) / s_pos

    # neg
    diag = torch.stack([rotmat[:, 0, 0], rotmat[:, 1, 1], rotmat[:, 2, 2]])
    max_val, max_ind = torch.max(diag, dim=0)

    s[~flag] = 2 * torch.sqrt(1 - trace[~flag] + 2 * max_val[~flag]) + 1e-16

    f0 = ~flag * (max_ind == 0)
    s0 = s[f0]
    quaternion[f0, 0] = (rotmat[f0, 2, 1] - rotmat[f0, 1, 2]) / s0
    quaternion[f0, 1] = s0 / 4
    quaternion[f0, 2] = (rotmat[f0, 0, 1] + rotmat[f0, 1, 0]) / s0
    quaternion[f0, 3] = (rotmat[f0, 0, 2] + rotmat[f0, 2, 0]) / s0

    f1 = ~flag * (max_ind == 1)
    s1 = s[f1]
    quaternion[f1, 0] = (rotmat[f1, 0, 2] - rotmat[f1, 2, 0]) / s1
    quaternion[f1, 1] = (rotmat[f1, 0, 1] + rotmat[f1, 1, 0]) / s1
    quaternion[f1, 2] = s1 / 4
    quaternion[f1, 3] = (rotmat[f1, 1, 2] + rotmat[f1, 2, 1]) / s1

    f2 = ~flag * (max_ind == 2)
    s2 = s[f2]
    quaternion[f2, 0] = (rotmat[f2, 1, 0] - rotmat[f2, 0, 1]) / s2
    quaternion[f2, 1] = (rotmat[f2, 0, 2] + rotmat[f2, 2, 0]) / s2
    quaternion[f2, 2] = (rotmat[f2, 1, 2] + rotmat[f2, 2, 1]) / s2
    quaternion[f2, 3] = s2 / 4

    return quaternion