import torch


def angles_to_rotation_matrix(yaw, pitch, roll, degrees=True):
    """ Convert pitch, yaw, roll angles to a right-handed rotation matrix. It can be used to post-multiply
        row vectors: rotated = vector * R.

        Based on RotationMatrix.m from the AFLW2000 and 300W LP datasets.

        :param pitch pitch (x) angle
        :param yaw yaw (y) angle
        :param roll roll (z) angle
        :param degrees if True, the angles are in degrees, otherwise in radians.
    """
    # Negate to make the matrix consistent with the head position.
    yaw = -yaw
    roll = -roll

    if degrees:
        pitch = torch.deg2rad(pitch)
        yaw = torch.deg2rad(yaw)
        roll = torch.deg2rad(roll)

    rx = torch.tensor([[1, 0, 0], [0, torch.cos(pitch), torch.sin(pitch)], [0, -torch.sin(pitch), torch.cos(pitch)]])
    ry = torch.tensor([[torch.cos(yaw), 0, -torch.sin(yaw)], [0, 1, 0], [torch.sin(yaw), 0, torch.cos(yaw)]])
    rz = torch.tensor([[torch.cos(roll), torch.sin(roll), 0], [-torch.sin(roll), torch.cos(roll), 0], [0, 0, 1]])

    r = torch.matmul(torch.matmul(rx, ry), rz)
    r = torch.transpose(r, 0, 1)
    return r