text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from typing import Any, Dict, List, Optional
  
  import magnum as mn
  
  import habitat_sim
  from habitat.sims.habitat_simulator.debug_visualizer import DebugVisualizer
  
  
  def register_custom_wireframe_box_template(
      sim: habitat_sim.Simulator,
      size: mn.Vector3,
      template_name: str = "custom_wireframe_box",
  ) -> str:
      """
      Generate and register a custom template for a wireframe box of given size.
      Return the new template's handle.
      """
      obj_attr_mgr = sim.get_object_template_manager()
      cube_template = obj_attr_mgr.get_template_by_handle(
          obj_attr_mgr.get_template_handles("cubeWireframe")[0]
      )
      cube_template.scale = size
      obj_attr_mgr.register_template(cube_template, template_name)
      return template_name
  
  
  def add_wire_box(
      sim: habitat_sim.Simulator,
      size: mn.Vector3,
      center: mn.Vector3,
      attach_to: Optional[habitat_sim.scene.SceneNode] = None,
      orientation: Optional[mn.Quaternion] = None,
  ) -> habitat_sim.physics.ManagedRigidObject:
      """
      Generate a wire box object and optionally attach it to another existing object (automatically applies object scale).
      Returns the new object.
      """
      if orientation is None:
          orientation = mn.Quaternion()
      box_template_handle = register_custom_wireframe_box_template(sim, size)
      new_object = sim.get_rigid_object_manager().add_object_by_template_handle(
          box_template_handle, attach_to
      )
      new_object.motion_type = habitat_sim.physics.MotionType.KINEMATIC
      new_object.collidable = False
      # translate to local offset if attached or global offset if not
      new_object.translation = center
      new_object.rotation = orientation
      return new_object
  
  
  def add_transformed_wire_box(
      sim: habitat_sim.Simulator,
      size: mn.Vector3,
      transform: Optional[mn.Matrix4] = None,
  ) -> habitat_sim.physics.ManagedRigidObject:
      """
      Generate a transformed wire box in world space.
      Returns the new object.
      """
      if transform is None:
          transform = mn.Matrix4()
      box_template_handle = register_custom_wireframe_box_template(sim, size)
      new_object = sim.get_rigid_object_manager().add_object_by_template_handle(
          box_template_handle
      )
      new_object.motion_type = habitat_sim.physics.MotionType.KINEMATIC
      new_object.collidable = False
      # translate to local offset if attached or global offset if not
      new_object.transformation = transform
      return new_object
  
  
  def add_viz_sphere(
      sim: habitat_sim.Simulator, radius: float, pos: mn.Vector3
  ) -> habitat_sim.physics.ManagedRigidObject:
      """
      Add a visualization-only sphere to the world at a global position.
      Returns the new object.
      """
      obj_attr_mgr = sim.get_object_template_manager()
      sphere_template = obj_attr_mgr.get_template_by_handle(
          obj_attr_mgr.get_template_handles("icosphereWireframe")[0]
      )
      sphere_template.scale = mn.Vector3(radius)
      obj_attr_mgr.register_template(sphere_template, "viz_sphere")
      new_object = sim.get_rigid_object_manager().add_object_by_template_handle(
          "viz_sphere"
      )
      new_object.motion_type = habitat_sim.physics.MotionType.KINEMATIC
      new_object.collidable = False
      new_object.translation = pos
      return new_object
  
  
  def get_bb_corners(
      obj: habitat_sim.physics.ManagedRigidObject,
  ) -> List[mn.Vector3]:
      """
      Return a list of object bounding box corners in object local space.
      """
      bb = obj.root_scene_node.cumulative_bb
      return [
          bb.back_bottom_left,
          bb.back_bottom_right,
          bb.back_top_right,
          bb.back_top_left,
          bb.front_top_left,
          bb.front_top_right,
          bb.front_bottom_right,
          bb.front_bottom_left,
      ]
  
  
  def get_ao_global_bb(
      obj: habitat_sim.physics.ManagedArticulatedObject,
  ) -> mn.Range3D:
      """
      Compute the cumulative bounding box of an ArticulatedObject by merging all link bounding boxes.
      """
      cumulative_global_bb = mn.Range3D()
      for link_ix in range(-1, obj.num_links):
          link_node = obj.get_link_scene_node(link_ix)
          bb = link_node.cumulative_bb
          global_bb = habitat_sim.geo.get_transformed_bb(
              bb, link_node.transformation
          )
          cumulative_global_bb = mn.math.join(cumulative_global_bb, global_bb)
      return cumulative_global_bb
  
  
  def bb_ray_prescreen(
      sim: habitat_sim.Simulator,
      obj: habitat_sim.physics.ManagedRigidObject,
      support_obj_ids: Optional[List[int]] = None,
      check_all_corners: bool = False,
  ) -> Dict[str, Any]:
      """
      Pre-screen a potential placement by casting rays in the gravity direction from the object center of mass (and optionally each corner of its bounding box) checking for interferring objects below.
  
      :param sim: The Simulator instance.
      :param obj: The RigidObject instance.
      :param support_obj_ids: A list of object ids designated as valid support surfaces for object placement. Contact with other objects is a criteria for placement rejection.
      :param check_all_corners: Optionally cast rays from all bounding box corners instead of only casting a ray from the center of mass.
      """
      if support_obj_ids is None:
          # set default support surface to stage/ground mesh
          support_obj_ids = [-1]
      lowest_key_point: mn.Vector3 = None
      lowest_key_point_height = None
      highest_support_impact: mn.Vector3 = None
      highest_support_impact_height = None
      highest_support_impact_with_stage = False
      raycast_results = []
      gravity_dir = sim.get_gravity().normalized()
      object_local_to_global = obj.transformation
      bb_corners = get_bb_corners(obj)
      key_points = [mn.Vector3(0)] + bb_corners  # [COM, c0, c1 ...]
      support_impacts: Dict[int, mn.Vector3] = {}  # indexed by keypoints
      for ix, key_point in enumerate(key_points):
          world_point = object_local_to_global.transform_point(key_point)
          # NOTE: instead of explicit Y coordinate, we project onto any gravity vector
          world_point_height = world_point.projected_onto_normalized(
              -gravity_dir
          ).length()
          if (
              lowest_key_point is None
              or lowest_key_point_height > world_point_height
          ):
              lowest_key_point = world_point
              lowest_key_point_height = world_point_height
          # cast a ray in gravity direction
          if ix == 0 or check_all_corners:
              ray = habitat_sim.geo.Ray(world_point, gravity_dir)
              raycast_results.append(sim.cast_ray(ray))
              # classify any obstructions before hitting the support surface
              for hit in raycast_results[-1].hits:
                  if hit.object_id == obj.object_id:
                      continue
                  elif hit.object_id in support_obj_ids:
                      hit_point = ray.origin + ray.direction * hit.ray_distance
                      support_impacts[ix] = hit_point
                      support_impact_height = mn.math.dot(
                          hit_point, -gravity_dir
                      )
  
                      if (
                          highest_support_impact is None
                          or highest_support_impact_height
                          < support_impact_height
                      ):
                          highest_support_impact = hit_point
                          highest_support_impact_height = support_impact_height
                          highest_support_impact_with_stage = hit.object_id == -1
  
                  # terminates at the first non-self ray hit
                  break
      # compute the relative base height of the object from its lowest bb corner and COM
      base_rel_height = (
          lowest_key_point_height
          - obj.translation.projected_onto_normalized(-gravity_dir).length()
      )
  
      # account for the affects of stage mesh margin
      margin_offset = (
          0
          if not highest_support_impact_with_stage
          else sim.get_stage_initialization_template().margin
      )
  
      surface_snap_point = (
          None
          if 0 not in support_impacts
          else support_impacts[0]
          + gravity_dir * (base_rel_height - margin_offset)
      )
      # return list of obstructed and grounded rays, relative base height, distance to first surface impact, and ray results details
      return {
          "base_rel_height": base_rel_height,
          "surface_snap_point": surface_snap_point,
          "raycast_results": raycast_results,
      }
  
  
  def snap_down(
      sim: habitat_sim.Simulator,
      obj: habitat_sim.physics.ManagedRigidObject,
      support_obj_ids: Optional[List[int]] = None,
      vdb: Optional[DebugVisualizer] = None,
  ) -> bool:
      """
      Attempt to project an object in the gravity direction onto the surface below it.
  
      :param sim: The Simulator instance.
      :param obj: The RigidObject instance.
      :param support_obj_ids: A list of object ids designated as valid support surfaces for object placement. Contact with other objects is a criteria for placement rejection. If none provided, default support surface is the stage/ground mesh (-1).
      :param vdb: Optionally provide a DebugVisualizer (vdb) to render debug images of each object's computed snap position before collision culling.
  
      Reject invalid placements by checking for penetration with other existing objects.
      Returns boolean success.
      If placement is successful, the object state is updated to the snapped location.
      If placement is rejected, object position is not modified and False is returned.
  
      To use this utility, generate an initial placement for any object above any of the designated support surfaces and call this function to attempt to snap it onto the nearest surface in the gravity direction.
      """
      cached_position = obj.translation
  
      if support_obj_ids is None:
          # set default support surface to stage/ground mesh
          support_obj_ids = [-1]
  
      bb_ray_prescreen_results = bb_ray_prescreen(sim, obj, support_obj_ids)
  
      if bb_ray_prescreen_results["surface_snap_point"] is None:
          # no support under this object, return failure
          return False
  
      # finish up
      if bb_ray_prescreen_results["surface_snap_point"] is not None:
          # accept the final location if a valid location exists
          obj.translation = bb_ray_prescreen_results["surface_snap_point"]
          if vdb is not None:
              vdb.get_observation(obj.translation)
          sim.perform_discrete_collision_detection()
          cps = sim.get_physics_contact_points()
          for cp in cps:
              if (
                  cp.object_id_a == obj.object_id
                  or cp.object_id_b == obj.object_id
              ) and (
                  (cp.contact_distance < -0.01)
                  or not (
                      cp.object_id_a in support_obj_ids
                      or cp.object_id_b in support_obj_ids
                  )
              ):
                  obj.translation = cached_position
                  # print(f" Failure: contact in final position w/ distance = {cp.contact_distance}.")
                  # print(f" Failure: contact in final position with non support object {cp.object_id_a} or {cp.object_id_b}.")
                  return False
          return True
      else:
          # no valid position found, reset and return failure
          obj.translation = cached_position
          return False
  
  
  def get_all_object_ids(sim: habitat_sim.Simulator) -> Dict[int, str]:
      """
      Generate a dict mapping all active object ids to a descriptive string containing the object instance handle and, for ArticulatedLinks, the link name.
      """
      rom = sim.get_rigid_object_manager()
      aom = sim.get_articulated_object_manager()
  
      object_id_map = {}
  
      for _object_handle, rigid_object in rom.get_objects_by_handle_substring(
          ""
      ).items():
          object_id_map[rigid_object.object_id] = rigid_object.handle
  
      for _object_handle, ao in aom.get_objects_by_handle_substring("").items():
          object_id_map[ao.object_id] = ao.handle
          for object_id, link_ix in ao.link_object_ids.items():
              object_id_map[object_id] = (
                  ao.handle + " -- " + ao.get_link_name(link_ix)
              )
  
      return object_id_map
  