text: # coding=utf-8
  # Copyright 2022 The Reach ML Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Assortment of utilities to interact with bullet within g3."""
  import dataclasses
  import datetime
  import getpass
  import gzip
  import json
  import os
  import time
  from typing import Any, Dict, List, Optional, Tuple
  
  from absl import logging
  from ibc.environments.utils.pose3d import Pose3d
  import numpy as np
  from scipy.spatial import transform
  import six
  import tensorflow as tf
  
  
  import pybullet
  import pybullet_data
  import pybullet_utils.bullet_client as bullet_client
  
  Vec3 = Tuple[float, float, float]
  Vec4 = Tuple[float, float, float, float]
  PYBULLET_STATE_VERSION = 2  # Basic versioning of serialized pybullet state.
  
  
  # Note about rotation_to_matrix and matrix_to_rotation below:
  # The abstractions below allow us to use older versions of scipy.
  def rotation_to_matrix(rotation):
    if hasattr(rotation, 'as_dcm'):
      return rotation.as_dcm()
    else:
      assert hasattr(rotation, 'as_matrix')
      return rotation.as_matrix()
  
  
  def matrix_to_rotation(matrix):
    if hasattr(transform.Rotation, 'from_dcm'):
      return transform.Rotation.from_dcm(matrix)
    else:
      assert hasattr(transform.Rotation, 'from_matrix')
      return transform.Rotation.from_matrix(matrix)
  
  
  def load_urdf(pybullet_client, file_path, *args, **kwargs):
    """Loads the given URDF filepath."""
  
    # Handles most general file open case.
    try:
      if os.path.exists(file_path):
        return pybullet_client.loadURDF(file_path, *args, **kwargs)
    except pybullet_client.error:
      pass
  
  
    try:
      if file_path.startswith('third_party/py/ibc/'):
        file_path = file_path[len('third_party/py/'):]
      if file_path.startswith(
          'third_party/bullet/examples/pybullet/gym/pybullet_data/'):
        pybullet_client.setAdditionalSearchPath(pybullet_data.getDataPath())
        file_path = file_path[55:]
      elif file_path.startswith('robotics/'):
        pybullet_client.setAdditionalSearchPath(os.environ['PYTHONPATH'])
        file_path = file_path[9:]
      elif file_path.startswith('third_party/'):
        pybullet_client.setAdditionalSearchPath(os.environ['PYTHONPATH'])
  
      logging.info('Loading URDF %s', file_path)
      return pybullet_client.loadURDF(file_path, *args, **kwargs)
    except pybullet.error:
      raise FileNotFoundError('Cannot load the URDF file {}'.format(file_path))
  
  
  
  
  def add_visual_sphere(client, center=(0, 0, 0), radius=0.1,
                        rgba=(0.5, 0.5, 0.5, 0.5)):
    """Add a sphere to bullet scene (visual only, no physics).
  
    Args:
      client: pybullet client (or pybullet library handle).
      center: Center of sphere.
      radius: Sphere radius.
      rgba: rgba color of sphere.
  
    Returns:
      Unique integer bullet id of constructed object.
    """
    vis_obj_id = client.createVisualShape(
        client.GEOM_SPHERE, radius=radius, rgbaColor=rgba)
    obj_id = client.createMultiBody(
        baseCollisionShapeIndex=-1, baseVisualShapeIndex=vis_obj_id,
        basePosition=center)
    return obj_id
  
  
  def pybullet_mat_to_numpy_4x4(pybullet_matrix):
    assert len(pybullet_matrix) == 16, 'pybullet matrix should be len 16'
    return np.transpose(
        np.reshape(np.array(pybullet_matrix, dtype=np.float64), (4, 4)))
  
  
  def decompose_view_matrix(pybullet_view_matrix):
    """Decompose view matrix into pos + quat format (assumes mat is rigid!)."""
    # It would be MUCH better to use something from bullet, however pybullet does
    # not expose all of the linear algebra library.
    mat = pybullet_mat_to_numpy_4x4(pybullet_view_matrix)
  
    # View matrix is now:
    # | R_11 R_12 R_13 t_1 |
    # | R_21 R_22 R_23 t_2 |
    # | R_31 R_32 R_33 t_3 |
    # |    0    0    0   1 |
  
    # R is the inverse eye to target at orientation, and t is R * eye.
    mat_view_to_world = np.linalg.inv(mat)
  
    # mat_view_to_world is the view to world transform, therefore the translation
    # component of this matrix is simply the world space position (since mat *
    # (0, 0, 0, 1)) is just copying the right column.
    world_xyz_view = np.copy(mat_view_to_world[0:3, 3])
  
    mat_view_to_world[0:3, 3] = 0  # Zero out the position change.
    world_quat_view = matrix_to_rotation(mat_view_to_world).as_quat()
  
    return world_xyz_view, world_quat_view
  
  
  def world_obj_to_view(world_xyz_obj, world_quat_obj, camera_view, client):
    """Transform object into view space."""
    world_xyz_view, world_quat_view = decompose_view_matrix(camera_view)
    view_xyz_world, view_quat_world = client.invertTransform(
        world_xyz_view, world_quat_view)
    view_xyz_obj, view_quat_obj = client.multiplyTransforms(
        view_xyz_world, view_quat_world, world_xyz_obj, world_quat_obj)
  
    return view_xyz_obj, view_quat_obj
  
  
  def image_xy_to_view_ray(xy, cam_width, cam_height, proj_mat_inv):
    """Calculate view-space ray from pixel location."""
    # Recall (from http://www.songho.ca/opengl/gl_projectionmatrix.html):
    # xyzw_clip = M_proj * xyzw_eye, and
    # xyz_ndc = xyzw_clip[0:3] / xwzw_clip[3].
    xyz_ndc = np.array([2.0 * xy[0] / cam_width - 1.0,
                        -(2.0 * xy[1] / cam_height - 1.0),
                        0])  # in [-1, 1]
    xyzw_clip = np.concatenate([xyz_ndc, [1]])
    xyzw_eye = proj_mat_inv @ xyzw_clip
    origin = np.zeros(3)
    vec = xyzw_eye[:3] / max(np.linalg.norm(xyzw_eye[:3]), 1e-6)
    return origin, vec
  
  
  def view_ray_to_world_ray(origin, vec, view_mat_inv):
    """Transform view-space ray into world space."""
    origin = view_mat_inv @ np.concatenate([origin, [1]])
    vec = view_mat_inv @ np.concatenate([vec, [0]])
  
    return origin[:3], vec[:3]
  
  
  def ray_to_plane_test(ray_origin,
                        ray_vec,
                        plane_origin,
                        plane_normal):
    """Perform a ray-plane intersection test."""
    ln = np.dot(plane_normal, ray_vec)
    if abs(ln) < np.finfo(np.float32).eps:
      return None
  
    # Solve for the intersection fraction t.
    t = np.dot(plane_normal, plane_origin - ray_origin) / ln
    if t >= 0:
      return ray_origin + ray_vec * t
    else:
      return None
  
  
  def get_workspace(env):
    (workspace_origin,
     workspace_quat) = env.pybullet_client.getBasePositionAndOrientation(
         env.workspace_uid)
    workspace_normal = rotation_to_matrix(transform.Rotation.from_quat(
        workspace_quat))[2, 0:3]
  
    return workspace_origin, workspace_normal
  
  
  def reset_camera_pose(env, view_type):
    """Reset camera pose to canonical frame."""
    p = env.pybullet_client
  
    if view_type == 'POLICY':
      camera_info = p.getDebugVisualizerCamera()
      image_size = (camera_info[0], camera_info[1])
  
      viewm, _, front_position, lookat, _ = env.calc_camera_params(image_size)
      # TODO(tompson): Request Erwin add method to set view and projection matrix
      # of the debug visualizer directly (because using yaw and pitch is
      # inconsistent).
  
      euler = matrix_to_rotation(
          pybullet_mat_to_numpy_4x4(viewm)[0:3, 0:3]).as_euler(
              'xyz', degrees=False)
      pitch = euler[1]
      yaw = -euler[2]
      # The distance is a bit far away (the GL view has higher FOV).
      distance = np.linalg.norm(front_position - lookat) * 0.6
    elif view_type == 'TOP_DOWN':
      workspace_origin, _ = get_workspace(env)
      distance = 0.5
      lookat = workspace_origin
      yaw = np.pi / 2
      # Note: pi/2 pitch results in gimble lock and pybullet doesn't support it.
      pitch = - (np.pi / 2 - 1e-5)
    else:
      raise ValueError('unsupported view_type %s' % view_type)
    p.resetDebugVisualizerCamera(
        cameraDistance=distance,
        cameraYaw=360 * yaw / (2.0 * np.pi),
        cameraPitch=360 * pitch / (2.0 * np.pi),
        cameraTargetPosition=lookat)
  
  
  def _lists_to_tuple(obj):
    if isinstance(obj, list):
      return tuple([_lists_to_tuple(v) for v in obj])
    else:
      return obj
  
  
  @dataclasses.dataclass
  class ObjState():
    """A container for storing pybullet object state."""
  
    obj_id: int
  
    # base_pose: (xyz, quat).
    base_pose: Tuple[Vec3, Vec4]
    # base_vel: (vel, ang_vel).
    base_vel: Tuple[Vec3, Vec3]
    joint_info: Any
    joint_state: Any
  
    @staticmethod
    def get_bullet_state(client,
                         obj_id):
      """Read Pybullet internal state."""
      base_pose = client.getBasePositionAndOrientation(obj_id)
      base_vel = client.getBaseVelocity(obj_id)
  
      joint_info = []
      joint_state = []
      for i in range(client.getNumJoints(obj_id)):
        joint_state.append(client.getJointState(obj_id, i))
        joint_info.append(ObjState._get_joint_info(client, obj_id, i))
  
      return ObjState(obj_id=obj_id,
                      base_pose=base_pose,
                      base_vel=base_vel,
                      joint_info=tuple(joint_info),
                      joint_state=tuple(joint_state))
  
    @staticmethod
    def _get_joint_info(client, obj_id, joint_index):
      ji = client.getJointInfo(obj_id, joint_index)
      return tuple(
          [v if not isinstance(v, bytes) else v.decode('utf-8') for v in ji])
  
    def set_bullet_state(self, client, obj_id):
      """Hard set the current bullet state."""
      xyz, quat = self.base_pose
      client.resetBasePositionAndOrientation(obj_id, xyz, quat)
      vel, ang_vel = self.base_vel
      client.resetBaseVelocity(obj_id, vel, ang_vel)
  
      njoints = client.getNumJoints(obj_id)
      if njoints != len(self.joint_info) or njoints != len(self.joint_state):
        raise ValueError('Incorrect number of joint info state pairs.')
  
      for i, (joint_info, joint_state) in enumerate(
          zip(self.joint_info, self.joint_state)):
        joint_index = joint_info[0]
        if joint_index != i:
          raise ValueError('Joint index mismatch.')
  
        # Check that the current joint we're trying to restore state for has the
        # same info as the state joint.
        cur_joint_info = ObjState._get_joint_info(client, obj_id, joint_index)
        if cur_joint_info != joint_info:
          raise ValueError('joint_info mismatch %s vs %s (expected)' % (
              str(cur_joint_info), str(joint_info)))
        joint_position = joint_state[0]
        joint_velocity = joint_state[1]
        client.resetJointState(
            obj_id, i, targetValue=joint_position, targetVelocity=joint_velocity)
  
    def serialize(self):
      return {'obj_id': self.obj_id,
              'base_pose': self.base_pose,
              'base_vel': self.base_vel,
              'joint_info': self.joint_info,
              'joint_state': self.joint_state}
  
    @staticmethod
    def deserialize(data):
      return ObjState(obj_id=_lists_to_tuple(data['obj_id']),
                      base_pose=_lists_to_tuple(data['base_pose']),
                      base_vel=_lists_to_tuple(data['base_vel']),
                      joint_info=_lists_to_tuple(data['joint_info']),
                      joint_state=_lists_to_tuple(data['joint_state']))
  
  
  @dataclasses.dataclass
  class XarmState(ObjState):
    """A container for storing pybullet robot state."""
  
    # The set point of the robot's controller.
    target_effector_pose: Pose3d
    goal_translation: Optional[Vec3]
  
    @staticmethod
    def get_bullet_state(client,
                         obj_id,
                         target_effector_pose,
                         goal_translation):
      if goal_translation is not None:
        goal_translation = tuple(goal_translation.tolist())
      return XarmState(
          **dataclasses.asdict(ObjState.get_bullet_state(client, obj_id)),
          target_effector_pose=target_effector_pose,
          goal_translation=goal_translation)
  
    def serialize(self):
      data = ObjState.serialize(self)
      data['target_effector_pose'] = self.target_effector_pose.serialize()
      if self.goal_translation is not None:
        data['goal_translation'] = self.goal_translation
      else:
        data['goal_translation'] = []
      return data
  
    @staticmethod
    def deserialize(data):
      goal_translation = (None if not data['goal_translation'] else
                          _lists_to_tuple(data['goal_translation']))
      return XarmState(
          obj_id=data['obj_id'],
          base_pose=_lists_to_tuple(data['base_pose']),
          base_vel=_lists_to_tuple(data['base_vel']),
          joint_info=_lists_to_tuple(data['joint_info']),
          joint_state=_lists_to_tuple(data['joint_state']),
          goal_translation=goal_translation,
          target_effector_pose=Pose3d.deserialize(data['target_effector_pose']))
  
  
  def _serialize_pybullet_state(pybullet_state):
    """Convert data to POD types."""
    if isinstance(pybullet_state, list):
      return [_serialize_pybullet_state(entry) for entry in pybullet_state]
    elif isinstance(pybullet_state, dict):
      assert '_serialized_obj_name' not in pybullet_state
      return {key: _serialize_pybullet_state(value)
              for key, value in pybullet_state.items()}
    elif isinstance(pybullet_state, (XarmState, ObjState)):
      return {
          '_serialized_obj_name': type(pybullet_state).__name__,
          '_serialized_data': pybullet_state.serialize()
      }
    elif isinstance(pybullet_state, int):
      return pybullet_state
    else:
      raise ValueError('Unhandled type for object %s, type %s' %
                       (str(pybullet_state), type(pybullet_state)))
  
  
  def _deserialize_pybullet_state(state):
    """Parse data from POD types."""
    if isinstance(state, list):
      return [_deserialize_pybullet_state(item) for item in state]
    elif isinstance(state, dict):
      if '_serialized_obj_name' in state:
        if state['_serialized_obj_name'] == XarmState.__name__:
          return XarmState.deserialize(state['_serialized_data'])
        elif state['_serialized_obj_name'] == ObjState.__name__:
          return ObjState.deserialize(state['_serialized_data'])
        else:
          raise ValueError('Unsupported: %s' % state['_serialized_obj_name'])
      else:
        return {key: _deserialize_pybullet_state(value)
                for key, value in state.items()}
    elif isinstance(state, int):
      return state
    else:
      raise ValueError('Unhandled type for object %s' % str(state))
  
  
  def write_pybullet_state(filename, pybullet_state, task,
                           actions = None):
    """Serialize pybullet state to json file."""
    data = {
        'pybullet_state': _serialize_pybullet_state(pybullet_state),
        'state_version': PYBULLET_STATE_VERSION,
        'ts_ms': int(time.mktime(datetime.datetime.now().timetuple())) * 1000,
        'user': getpass.getuser(),
        'task': task,
        'actions': actions if actions is not None else [],
    }
    with tf.io.gfile.GFile(filename, 'wb') as fdata:
      with gzip.GzipFile(fileobj=fdata) as fzip:
        fzip.write(json.dumps(data).encode('utf-8'))
  
  
  def read_pybullet_state(filename):
    """Deserialize pybullet state from json file."""
    with tf.io.gfile.GFile(filename, 'rb') as fdata:
      with gzip.GzipFile(fileobj=fdata) as fzip:
        data = json.loads(fzip.read().decode('utf-8'))
  
    assert isinstance(data, dict)
  
    if data['state_version'] != PYBULLET_STATE_VERSION:
      raise ValueError('incompatible state data (version %d, expected %d)' % (
          data['state_version'], PYBULLET_STATE_VERSION))
  
    data['pybullet_state'] = _deserialize_pybullet_state(data['pybullet_state'])
    return data
  