text: # Copyright (c) 2018-present, Facebook, Inc.
  # All rights reserved.
  #
  # From
  # https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py
  # 
  # Attribution non-commerical license available at 
  # https://github.com/facebookresearch/QuaterNet
  
  import torch
  import numpy as np
  
  # PyTorch-backed implementations
  
  def qmul(q, r):
      """
      Multiply quaternion(s) q with quaternion(s) r.
      Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
      Returns q*r as a tensor of shape (*, 4).
      """
      assert q.shape[-1] == 4
      assert r.shape[-1] == 4
      
      original_shape = q.shape
      
      # Compute outer product
      terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
  
      w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
      x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
      y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
      z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
      return torch.stack((w, x, y, z), dim=1).view(original_shape)
  
  def qrot(q, v):
      """
      Rotate vector(s) v about the rotation described by quaternion(s) q.
      Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
      where * denotes any number of dimensions.
      Returns a tensor of shape (*, 3).
      """
      assert q.shape[-1] == 4
      assert v.shape[-1] == 3
      assert q.shape[:-1] == v.shape[:-1]
      
      original_shape = list(v.shape)
      q = q.view(-1, 4)
      v = v.view(-1, 3)
      
      qvec = q[:, 1:]
      uv = torch.cross(qvec, v, dim=1)
      uuv = torch.cross(qvec, uv, dim=1)
      return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
  
  def qeuler(q, order, epsilon=0):
      """
      Convert quaternion(s) q to Euler angles.
      Expects a tensor of shape (*, 4), where * denotes any number of dimensions.
      Returns a tensor of shape (*, 3).
      """
      assert q.shape[-1] == 4
      
      original_shape = list(q.shape)
      original_shape[-1] = 3
      q = q.view(-1, 4)
      
      q0 = q[:, 0]
      q1 = q[:, 1]
      q2 = q[:, 2]
      q3 = q[:, 3]
      
      if order == 'xyz':
          x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2*(q1 * q1 + q2 * q2))
          y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1+epsilon, 1-epsilon))
          z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2*(q2 * q2 + q3 * q3))
      elif order == 'yzx':
          x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2*(q1 * q1 + q3 * q3))
          y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2*(q2 * q2 + q3 * q3))
          z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1+epsilon, 1-epsilon))
      elif order == 'zxy':
          x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1+epsilon, 1-epsilon))
          y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2*(q1 * q1 + q2 * q2))
          z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2*(q1 * q1 + q3 * q3))
      elif order == 'xzy':
          x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2*(q1 * q1 + q3 * q3))
          y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2*(q2 * q2 + q3 * q3))
          z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1+epsilon, 1-epsilon))
      elif order == 'yxz':
          x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1+epsilon, 1-epsilon))
          y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2*(q1 * q1 + q2 * q2))
          z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2*(q1 * q1 + q3 * q3))
      elif order == 'zyx':
          x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2*(q1 * q1 + q2 * q2))
          y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1+epsilon, 1-epsilon))
          z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2*(q2 * q2 + q3 * q3))
      else:
          raise
  
      return torch.stack((x, y, z), dim=1).view(original_shape)
  
  # Numpy-backed implementations
  
  def qmul_np(q, r):
      q = torch.from_numpy(q).contiguous()
      r = torch.from_numpy(r).contiguous()
      return qmul(q, r).numpy()
  
  def qrot_np(q, v):
      q = torch.from_numpy(q).contiguous()
      v = torch.from_numpy(v).contiguous()
      return qrot(q, v).numpy()
  
  def qeuler_np(q, order, epsilon=0, use_gpu=False):
      if use_gpu:
          q = torch.from_numpy(q).cuda()
          return qeuler(q, order, epsilon).cpu().numpy()
      else:
          q = torch.from_numpy(q).contiguous()
          return qeuler(q, order, epsilon).numpy()
  
  def qfix(q):
      """
      Enforce quaternion continuity across the time dimension by selecting
      the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)
      between two consecutive frames.
      
      Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.
      Returns a tensor of the same shape.
      """
      assert len(q.shape) == 3
      assert q.shape[-1] == 4
      
      result = q.copy()
      dot_products = np.sum(q[1:]*q[:-1], axis=2)
      mask = dot_products < 0
      mask = (np.cumsum(mask, axis=0)%2).astype(bool)
      result[1:][mask] *= -1
      return result
  
  def expmap_to_quaternion(e):
      """
      Convert axis-angle rotations (aka exponential maps) to quaternions.
      Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
      Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
      Returns a tensor of shape (*, 4).
      """
      assert e.shape[-1] == 3
      
      original_shape = list(e.shape)
      original_shape[-1] = 4
      e = e.reshape(-1, 3)
  
      theta = torch.norm(e, dim=1).reshape(-1, 1)
      w = torch.cos(0.5*theta).reshape(-1, 1)
      xyz = torch.sin(0.5*theta)* (e / theta)
      return torch.cat((w, xyz), dim=1).reshape(original_shape)
  
  def euler_to_quaternion(e, order):
      """
      Convert Euler angles to quaternions.
      """
      assert e.shape[-1] == 3
      
      original_shape = list(e.shape)
      original_shape[-1] = 4
      
      e = e.reshape(-1, 3)
      
      x = e[:, 0]
      y = e[:, 1]
      z = e[:, 2]
      
      rx = torch.stack((torch.cos(x/2), torch.sin(x/2), torch.zeros_like(x), torch.zeros_like(x)), axis=1)
      ry = torch.stack((torch.cos(y/2), torch.zeros_like(y), torch.sin(y/2), torch.zeros_like(y)), axis=1)
      rz = torch.stack((torch.cos(z/2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z/2)), axis=1)
  
      result = None
      for coord in order:
          if coord == 'x':
              r = rx
          elif coord == 'y':
              r = ry
          elif coord == 'z':
              r = rz
          else:
              raise
          if result is None:
              result = r
          else:
              result = qmul(result, r)
              
      # Reverse antipodal representation to have a non-negative "w"
      if order in ['xyz', 'yzx', 'zxy']:
          result *= -1
      
      return result.reshape(original_shape)
  
  def quat2mat(quat):
      """
      From https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py
      Convert quaternion coefficients to rotation matrix.
      Args:
          quat: first three coeff of quaternion of rotation. fourht is then computed to have a norm of 1 -- size = [B, 3]
      Returns:
          Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
      """
      norm_quat = quat #torch.cat([quat[:,:1].detach()*0 + 1, quat], dim=1)
      norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)
      w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]
  
      B = quat.size(0)
  
      w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
      wx, wy, wz = w*x, w*y, w*z
      xy, xz, yz = x*y, x*z, y*z
  
      rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
                            2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
                            2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).reshape(B, 3, 3)
      return rotMat
  
  def rotation_matrix_to_quaternion(
          rotation_matrix: torch.Tensor,
          eps: float = 1e-8) -> torch.Tensor:
      r"""
  
      From https://github.com/kornia/kornia/blob/master/kornia/geometry/conversions.py
  
      Convert 3x3 rotation matrix to 4d quaternion vector.
      The quaternion vector has components in (x, y, z, w) format.
      Args:
          rotation_matrix (torch.Tensor): the rotation matrix to convert.
          eps (float): small value to avoid zero division. Default: 1e-8.
      Return:
          torch.Tensor: the rotation in quaternion.
      Shape:
          - Input: :math:`(*, 3, 3)`
          - Output: :math:`(*, 4)`
      Example:
          >>> input = torch.rand(4, 3, 3)  # Nx3x3
          >>> output = kornia.rotation_matrix_to_quaternion(input)  # Nx4
      """
      if not isinstance(rotation_matrix, torch.Tensor):
          raise TypeError("Input type is not a torch.Tensor. Got {}".format(
              type(rotation_matrix)))
  
      if not rotation_matrix.shape[-2:] == (3, 3):
          raise ValueError(
              "Input size must be a (*, 3, 3) tensor. Got {}".format(
                  rotation_matrix.shape))
  
      def safe_zero_division(numerator: torch.Tensor,
                             denominator: torch.Tensor) -> torch.Tensor:
          eps: float = torch.finfo(numerator.dtype).tiny  # type: ignore
          return numerator / torch.clamp(denominator, min=eps)
  
      rotation_matrix_vec: torch.Tensor = rotation_matrix.reshape(
          *rotation_matrix.shape[:-2], 9)
  
      m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.chunk(
          rotation_matrix_vec, chunks=9, dim=-1)
  
      trace: torch.Tensor = m00 + m11 + m22
  
      def trace_positive_cond():
          sq = torch.sqrt(trace + 1.0) * 2.  # sq = 4 * qw.
          qw = 0.25 * sq
          qx = safe_zero_division(m21 - m12, sq)
          qy = safe_zero_division(m02 - m20, sq)
          qz = safe_zero_division(m10 - m01, sq)
          return torch.cat([qx, qy, qz, qw], dim=-1)
  
      def cond_1():
          sq = torch.sqrt(1.0 + m00 - m11 - m22 + eps) * 2.  # sq = 4 * qx.
          qw = safe_zero_division(m21 - m12, sq)
          qx = 0.25 * sq
          qy = safe_zero_division(m01 + m10, sq)
          qz = safe_zero_division(m02 + m20, sq)
          return torch.cat([qx, qy, qz, qw], dim=-1)
  
      def cond_2():
          sq = torch.sqrt(1.0 + m11 - m00 - m22 + eps) * 2.  # sq = 4 * qy.
          qw = safe_zero_division(m02 - m20, sq)
          qx = safe_zero_division(m01 + m10, sq)
          qy = 0.25 * sq
          qz = safe_zero_division(m12 + m21, sq)
          return torch.cat([qx, qy, qz, qw], dim=-1)
  
      def cond_3():
          sq = torch.sqrt(1.0 + m22 - m00 - m11 + eps) * 2.  # sq = 4 * qz.
          qw = safe_zero_division(m10 - m01, sq)
          qx = safe_zero_division(m02 + m20, sq)
          qy = safe_zero_division(m12 + m21, sq)
          qz = 0.25 * sq
          return torch.cat([qx, qy, qz, qw], dim=-1)
  
      where_2 = torch.where(m11 > m22, cond_2(), cond_3())
      where_1 = torch.where(
          (m00 > m11) & (m00 > m22), cond_1(), where_2)
  
      quaternion: torch.Tensor = torch.where(
          trace > 0., trace_positive_cond(), where_1)
      return quaternion
  
  def transParamsToHomMatrix(q, t):
      """q = [V, 4], t = [V,3]"""
      N = q.size(0)
      R = quat2mat(q) # [V,3,3]
      Rt = torch.cat([R, t.unsqueeze(-1)], dim=2) # [V,3,4]
      hom_aug = torch.cat([torch.zeros([N, 1, 3]).double(),
                           torch.ones([N, 1, 1]).double()],
                          dim=2).to(Rt.device)
      RtHom = torch.cat([Rt, hom_aug], dim=1) # [V,4,4]
      return RtHom
  
  if __name__ == "__main__":
      print("Starting tests to try to find correspondence with Drake conventions")
  
      import numpy as np
      import drake
      from pydrake.common.eigen_geometry import Quaternion, AngleAxis, Isometry3
      from pydrake.math import RigidTransform, RollPitchYaw
  
      input_quat = np.array([ 8.5056e-06,  9.9844e-01,  5.9879e-02, -1.0553e-05])
      #input_quat = np.random.random(4)
      quat_norm = input_quat / np.linalg.norm(input_quat)
      drake_quat = Quaternion(quat_norm)
      drake_rpy = RollPitchYaw(drake_quat)
  
      torch_quat = torch.tensor(quat_norm)
      for order in ['xyz',
                    'yzx',
                    'zxy',
                    'xzy',
                    'yxz',
                    'zyx']:
          print("Order %s:" % order)
          torch_rpy = qeuler(torch.reshape(torch_quat, (1, 4)), order=order)
          print("\tDrake RPY:", drake_rpy.vector())
          print("\tTorch RPY:", torch_rpy.numpy())
          print("\"Back to quat: ", euler_to_quaternion(torch_rpy, order=order).numpy())
          print("\tOrig quat: ", torch_quat.numpy())
  
  