import torch


def mat2euler(mat):
    """ Convert rotation matrix to euler angles.
    
    https://github.com/pulkitag/pycaffe-utils/blob/master/rot_utils.py#L283

    Args:
        mat: rotation matrix in zyx format -- size = [B, 3, 3]
    Returns:
        angle: rotation angle along 3 axis (in radians, it's not unique) -- size = [B, 3]

    """
    cy_thresh = 1e-10
    cy = torch.sqrt(mat[:, 2, 2]*mat[:, 2, 2] + mat[:, 1, 2]*mat[:, 1, 2])
    
    if (cy > cy_thresh).any(): # cos(y) not close to zero, standard form
        z = torch.atan2(-mat[:, 0, 1],  mat[:, 0, 0]) # atan2(cos(y)*sin(z), cos(y)*cos(z))
        y = torch.atan2(mat[:, 0, 2],  cy) # atan2(sin(y), cy)
        x = torch.atan2(-mat[:, 1, 2], mat[:, 2, 2]) # atan2(cos(y)*sin(x), cos(x)*cos(y))
    else: # cos(y) (close to) zero, so x -> 0.0 (see above)
        # so r21 -> sin(z), r22 -> cos(z) and
        z = torch.atan2(mat[:, 1, 0],  mat[:, 1, 1])
        y = torch.atan2(mat[:, 0, 2],  cy) # atan2(sin(y), cy)
        x = torch.zeros_like(mat[:, 0, 0])

    return torch.cat([x.unsqueeze(-1), y.unsqueeze(-1), z.unsqueeze(-1)], -1).view(-1, 3)