import torch


def quat_to_rotMat(q):
    """Convert quaternions to rotation matrices.

    Using equation provided in XSens MVN Manual:
    https://www.xsens.com/hubfs/Downloads/usermanual/MVN_User_Manual.pdf

    Args:
        q (torch.Tensor): quaternion(s) to convert to rotation matrix format

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

    assert q.shape[1] == 4

    r0c0 = q[:, 0]**2 + q[:, 1]**2 - q[:, 2]**2 - q[:, 3]**2
    r0c1 = 2*q[:, 1]*q[:, 2] - 2*q[:, 0]*q[:, 3]
    r0c2 = 2*q[:, 1]*q[:, 3] + 2*q[:, 0]*q[:, 2]

    r1c0 = 2*q[:, 1]*q[:, 2] + 2*q[:, 0]*q[:, 3]
    r1c1 = q[:, 0]**2 - q[:, 1]**2 + q[:, 2]**2 - q[:, 3]**2
    r1c2 = 2*q[:, 2]*q[:, 3] - 2*q[:, 0]*q[:, 1]

    r2c0 = 2*q[:, 1]*q[:, 3] - 2*q[:, 0]*q[:, 2]
    r2c1 = 2*q[:, 2]*q[:, 3] + 2*q[:, 0]*q[:, 1]
    r2c2 = q[:, 0]**2 - q[:, 1]**2 - q[:, 2]**2 + q[:, 3]**2

    r0 = torch.stack([r0c0, r0c1, r0c2], dim=1)
    r1 = torch.stack([r1c0, r1c1, r1c2], dim=1)
    r2 = torch.stack([r2c0, r2c1, r2c2], dim=1)

    R = torch.stack([r0, r1, r2], dim=2)

    return R.permute(0, 2, 1)