import numpy as np

def rpy_to_matrix(roll, pitch, yaw, degrees=False):
    """
    URDF: R = Rz(yaw) @ Ry(pitch) @ Rx(roll)
    roll, pitch, yaw 默认为弧度；degrees=True 时按角度输入。
    """
    if degrees:
        roll, pitch, yaw = np.deg2rad([roll, pitch, yaw])

    cr, sr = np.cos(roll),  np.sin(roll)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw),   np.sin(yaw)

    Rx = np.array([[1, 0, 0],
                   [0, cr, -sr],
                   [0, sr,  cr]], dtype=float)
    Ry = np.array([[ cp, 0, sp],
                   [  0, 1,  0],
                   [-sp, 0, cp]], dtype=float)
    Rz = np.array([[cy, -sy, 0],
                   [sy,  cy, 0],
                   [ 0,   0, 1]], dtype=float)

    return Rz @ Ry @ Rx


def matrix_to_rpy(R, degrees=False, eps=1e-9):
    """
    从旋转矩阵恢复 URDF 的 (roll, pitch, yaw)，满足 R = Rz(yaw)·Ry(pitch)·Rx(roll)。
    返回弧度；degrees=True 时返回角度。
    含万向节锁处理。
    """
    R = np.asarray(R, dtype=float)
    assert R.shape == (3, 3)

    # 可选：小幅正交化以抑制数值误差（不想要可注释掉）
    # 使用极分解的简化：R ≈ U*V^T
    U, _, Vt = np.linalg.svd(R)
    R = U @ Vt

    # R[2,0] = -sin(pitch)
    sp_neg = R[2, 0]
    sp_neg = np.clip(sp_neg, -1.0, 1.0)

    # 非万向节锁：|sp_neg| < 1
    if abs(sp_neg) < 1.0 - eps:
        pitch = np.arcsin(-sp_neg)                 # pitch
        roll  = np.arctan2(R[2, 1], R[2, 2])       # roll
        yaw   = np.arctan2(R[1, 0], R[0, 0])       # yaw
    else:
        # 万向节锁：|sp_neg| ≈ 1，此时 yaw 与 roll 耦合
        # 令 yaw = 0，通过 R 的其它项恢复 roll
        pitch = np.pi/2 if sp_neg < 0 else -np.pi/2
        yaw   = 0.0
        # 当 pitch = ±pi/2 时，R[0,1] 与 R[1,1] 携带 roll 信息
        # 推导自 R = Rz(yaw)·Ry(±pi/2)·Rx(roll)
        roll = np.arctan2(-R[0, 1] if sp_neg < 0 else R[0, 1],
                          R[1, 1])

    if degrees:
        return tuple(np.rad2deg([roll, pitch, yaw]))
    return (roll, pitch, yaw)


# ---------- 简单自检 ----------
if __name__ == "__main__":
    # rpy = (0.3, -0.6, 1.2)  # roll, pitch, yaw (rad)
    # R = rpy_to_matrix(*rpy)
    # rpy_back = matrix_to_rpy(R)
    # print("R:\n", R)
    # print("rpy back:", rpy_back)
    # print("max abs diff:", np.max(np.abs(np.array(rpy) - np.array(rpy_back))))
    """
    3.54839600e-04, -9.96155522e-01,  8.76016580e-02,  1.49988098e-01,
          9.99999921e-01,  3.37579500e-04, -2.11843500e-04, -0.35,
          1.81456600e-04,  8.76017262e-02,  9.96155563e-01,
    """
    R = np.array([
        [3.54839600e-04, -9.96155522e-01, 8.76016580e-02],
        [9.99999921e-01, 3.37579500e-04, -2.11843500e-04],
        [1.81456600e-04, 8.76017262e-02, 9.96155563e-01]
    ])
    print(np.array(matrix_to_rpy(R)) * 180 / np.pi)