import torch


def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
    """
    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
    using Gram--Schmidt orthogonalization per Section B of [1].
    Args:
        d6: 6D rotation representation, of size (*, 6)

    Returns:
        batch of rotation matrices of size (*, 3, 3)

    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
    On the Continuity of Rotation Representations in Neural Networks.
    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
    Retrieved from http://arxiv.org/abs/1812.07035
    """

    a1, a2 = d6[..., :3], d6[..., 3:]
    b1 = torch.nn.functional.normalize(a1, dim=-1)
    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
    b2 = torch.nn.functional.normalize(b2, dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)
    return torch.stack((b1, b2, b3), dim=-2)


def action_to_pose_single(action, idx):

    idx_xyz = 9 * idx
    idx_rot = idx_xyz + 3

    xyz = action[:, idx_xyz:idx_xyz+3]
    rot = action[:, idx_rot:idx_rot+6]

    rot = rotation_6d_to_matrix(rot)
    pose = torch.eye(4).to(action.device)
    pose[:3, :3] = rot.clone()
    pose[:3, -1] = xyz.clone()
    pose = pose.unsqueeze(0)
    return pose


def action_to_pose(action):
    action0 = action_to_pose_single(action, 0)
    action1 = action_to_pose_single(action, 1)
    pose = torch.stack([action0,action1], 1)
    return pose