# srmp API Reference (condensed)

class Pose:
    p: NDArray  # [x, y, z]
    q: NDArray  # [w, x, y, z] quaternion
    def __init__(self, p=[0,0,0], q=[1,0,0,0])
    def __init__(self, matrix: NDArray)          # 4x4 transformation matrix
    def __mul__(self, other: Pose) -> Pose
    def __mul__(self, point: NDArray) -> NDArray  # transform a 3D point
    def inv(self) -> Pose
    def distance(self, other: Pose) -> float
    def to_transformation_matrix(self) -> NDArray  # 4x4

class GoalType:
    JOINTS: int
    POSITION: int
    POSE: int
    MULTI_GOAL_JOINTS: int
    MULTI_GOAL_POSITION: int
    MULTI_GOAL_POSE: int

class GoalConstraint:
    type: int
    joints: List[NDArray]     # used with JOINTS / MULTI_GOAL_JOINTS
    positions: List[NDArray]  # used with POSITION / MULTI_GOAL_POSITION, each [x, y, z]
    poses: List[Pose]         # used with POSE / MULTI_GOAL_POSE
    def __init__(self)
    def __init__(self, type: int)
    def __init__(self, type: int, joints: List[NDArray])      # for JOINTS / MULTI_GOAL_JOINTS
    def __init__(self, type: int, positions: List[NDArray])   # for POSITION / MULTI_GOAL_POSITION
    def __init__(self, type: int, poses: List[Pose])          # for POSE / MULTI_GOAL_POSE

class Trajectory:
    positions: List[List[float]]
    velocities: List[List[float]]
    accelerations: List[List[float]]
    times: List[float]

class PlannerInterface:
    def __init__(self, grid_config=None)  # occupancy-grid config for collision checking; leave default unless told otherwise
    # Robot management
    def add_robot(self, robot: str, name=None, srdf_path=None, end_effector=None, planned=True,
                  gravity=None, link_names=None, joint_names=None, gripper_joint_names=None) -> str
        # preferred: loads from registry by name, or from a URDF file path (srdf_path/end_effector
        # required for a path). Returns the actual assigned name (may differ if `name` collides).
        # gripper_joint_names overrides the registry entry's own gripper classification, if any.
    def add_articulation(self, name, end_effector, urdf_path, srdf_path="", link_names=[], joint_names=[],
                          gravity=[0,0,0], planned=True, gripper_joint_names=[]) -> str
        # manual URDF; returns assigned name
    def remove_articulation(self, name) -> None
    def set_base_pose(self, name, pose: Pose) -> None
    # Scene objects
    def add_box(self, name, size, pose=None) -> None
    def add_sphere(self, name, radius, pose=None) -> None
    def add_cylinder(self, name, radius, length, pose=None) -> None
    def add_mesh(self, name, filepath, pose=None, scale=[1,1,1]) -> None
    def add_point_cloud(self, name, points, colors=None, pose=None) -> None
    def remove_object(self, name) -> None
    def read_sim(self, sim, sim_type: str, articulations=None) -> None
        # bulk-import scene geometry (boxes/spheres/cylinders/meshes) from a running simulation.
        # sim_type: "sapien" | "genesis" | "swift" | "pybullet" | "mujoco" (swift raises NotImplementedError).
        # articulations: names of bodies/geoms to skip (e.g. robot links already added via add_robot).
    # Queries
    def get_articulation_names(self) -> List[str]
    def get_object_names(self) -> List[str]
    def get_active_joint_names(self, articulation_name) -> List[str]       # ALL actuated joints
    def get_move_group_joint_names(self, articulation_name) -> List[str]   # joints the planner controls (subset of active joints)
    def get_move_group_qpos_dim(self, articulation_name) -> int           # planning DOF count for this articulation
    def get_link_names(self, name) -> List[str]
    def get_link_index(self, name, link_name) -> int
    def get_move_group_joint_limits(self, articulation_name) -> NDArray  # (move_group_qpos_dim, 2) [lower, upper]; single-DOF joints only
    def get_gripper_joint_names(self, articulation_name) -> List[str]
    def has_articulation(self, name) -> bool
    def has_object(self, name) -> bool
    def is_articulation_planned(self, name) -> bool
    def set_articulation_planned(self, name, planned: bool) -> None
    def is_object_attached(self, name) -> bool
    # Kinematics
    def compute_fk(self, articulation_name, qpos) -> Pose
    def compute_ik(self, articulation_name, ee_pose, init_state_val) -> Tuple[bool, List[float]]
        # ee_pose accepts either a Pose or a flat List[float]
    def get_link_pose(self, articulation_name, link_name) -> Pose
    def get_qpos(self, articulation_name) -> NDArray  # current qpos, move-group order (same order set_qpos/plan use)
    def set_qpos(self, name, qpos) -> None
    def set_gripper_qpos(self, robot_name, gripper_qpos) -> None
    def set_qpos_all(self, state) -> None  # set qpos for all planned articulations at once
    def get_jacobian(self, articulation_name, link_name, local=False, move_group_only=True) -> NDArray
        # (6, N) Jacobian at the articulation's *current* qpos. local: link frame vs world frame.
        # move_group_only: only move-group columns (default) vs all model DOF columns.
    def compute_jacobian(self, articulation_name, qpos, link_name, local=False, move_group_only=True) -> NDArray
        # same as get_jacobian but evaluated at an arbitrary move-group `qpos`, without mutating
        # the articulation's current state (non-move-group joints taken from current state)
    def compute_joint_velocities(self, articulation_name, link_name, twist, damping=0.0, local=False) -> NDArray
        # maps a desired 6D end-effector twist [linear; angular] to move-group joint velocities
        # via damped least squares (damping=0 -> plain pseudo-inverse; >0 more stable near singularities)
    def get_manipulability(self, articulation_name, link_name) -> float
        # Yoshikawa manipulability index sqrt(det(J J^T)) of the link's move-group Jacobian;
        # 0 at a kinematic singularity, higher = more dexterous
    # Collision
    def is_state_colliding(self, articulation_name="") -> bool  # empty name checks the whole scene
    def is_robot_colliding_with_objects(self, art_name) -> bool
    def distance_to_self_collision(self) -> float
    def distance_to_robot_collision(self) -> float
    def distance_to_collision(self) -> float
    def set_allowed_collision(self, name1, name2, allowed) -> None
    # Attachment
    def attach_object(self, object_name, robot_name, link_id, pose=None, is_link_name=False) -> None
    def detach_object(self, name, also_remove=False) -> bool
    def detach_all_objects(self, also_remove=False) -> bool
    def update_attached_bodies_pose(self) -> None  # refresh poses of all attached bodies from current robot state
    # Planning
    def make_planner(self, articulation_names: List[str], planner_context: Dict[str, str]) -> None
        # planner_context keys (all values are strings; names are case-sensitive):
        #   planner_id: single-agent "Astar"/"wAstar"/"ARAstar"/"MHAstar"/"wPASE",
        #               multi-agent (len(articulation_names) > 1) "ECBS"/"xECBS"
        #   heuristic_<articulation>, mprim_path_<articulation>  (per-agent, multi-agent only)
        #   weight_low_level_heuristic (default 1.0), high_level_focal_suboptimality (default 1.3),
        #   low_level_focal_suboptimality (default 1.3)
        #   verbose ("true"/"1" for search progress diagnostics, default "false")
    def plan(self, start, goal_constraint: GoalConstraint) -> Trajectory
    def plan_multi(self, start_states, goal_constraints) -> Dict[str, Trajectory]
    def plan_screw(self, articulation_name, link_name, start_qpos=None, end_pose: Pose = None,
                   axis_point=None, axis_direction=None, pitch=0.0, angle=None,
                   qpos_step=0.1, max_steps=10000) -> Trajectory
        # Cartesian screw-motion plan via Jacobian resolved-rate stepping (no IK solve); each
        # step is collision- and joint-limit-checked. Specify exactly one of:
        #   - end_pose: target pose for link_name (axis/pitch/angle derived automatically)
        #   - axis_direction + angle (+ optional axis_point, pitch): explicit screw motion,
        #     e.g. a hinge/door swing. axis_point defaults to link_name's current position.
        # start_qpos defaults to the articulation's current qpos; also sets it as a side effect.
        # Raises RuntimeError on inconsistent params, zero relative rotation, collision, joint
        # limit violation, singularity stall, or exceeding max_steps.
    def reset(self, reset_robots=True) -> None
    def print_available_planners(self) -> None  # prints available single/multi-agent planner IDs
    # Visualization
    def start_visualizer(self, type="viser", **kwargs)
        # Preferred way to get a live 3D viewer: constructs and attaches it, opens the
        # browser, and returns the visualizer instance. type: "viser" or "meshcat".
        # e.g. viz = planner.start_visualizer("viser", port=8080)
        # The returned object also exposes:
        #   viz.animate_trajectory(trajectories, dt=0.05, robot_name=None, fps=30.0) -> None
        #   viz.stop() -> None
    def get_visualizer(self, index=0)  # get an already-attached visualizer, or None
    def attach_visualizer(self, listener) -> None  # low-level: manually attach a VisualizerListener
    def detach_visualizer(self, listener) -> None  # low-level: manually detach a VisualizerListener

# Robot Registry
srmp.robots.list_available(fetch=True)  # {'remote': [...], 'local': [...], 'custom': [...]}
srmp.robots.download(name, force=False) -> Path       # downloads + caches; returns the robot's cache directory
srmp.robots.download_all(force=False) -> Path         # download every registry robot at once
srmp.robots.get(name) -> RobotInfo   # .urdf_path, .srdf_path, .end_effector, .description, .default_qpos, .joint_names
srmp.robots.info(name) -> RobotInfo  # alias for get()
srmp.robots.register(name, urdf_path, srdf_path, end_effector,
                      description=None, default_qpos=None, joint_names=None) -> None  # register a custom local robot
srmp.robots.unregister(name) -> None
srmp.robots.get_cache_dir() -> Path
srmp.robots.set_cache_dir(path) -> None
# Raises srmp.robots.RobotNotFoundError / srmp.robots.DownloadError on failure
