ardy/viz/viser_utils.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# ruff: noqa: I001
import os
import xml.etree.ElementTree as ET
from typing import List, Optional, Tuple, Union
import time

import numpy as np
import torch
import trimesh
import viser
import viser.transforms as tf

from scipy.interpolate import interp1d

from ardy.assets import skeleton_asset_path
from ardy.skeleton.kinematics import batch_rigid_transform
from ardy.skeleton import (
    CoreSkeleton27,
    G1Skeleton34,
    SOMASkeleton30,
    SOMASkeleton77,
    SkeletonBase,
)
from ardy.motion_rep.smooth_root import get_smooth_root_pos
from ardy.skeleton.transforms import global_rots_to_local_rots
from ardy.tools import to_numpy, to_torch
from ardy.viz.core_skin import CoreSkin
from ardy.viz.g1_rig import G1MeshRig
from ardy.viz.soma_skin import SOMASkin


_G1_JOINT_AXIS_INDEX_CACHE: Optional[dict[str, int]] = None


def _get_g1_joint_axis_indices() -> dict[str, int]:
    """Return a map from G1 joint names to a single rotation axis index."""
    global _G1_JOINT_AXIS_INDEX_CACHE
    if _G1_JOINT_AXIS_INDEX_CACHE is not None:
        return _G1_JOINT_AXIS_INDEX_CACHE

    xml_path = str(skeleton_asset_path("g1skel34", "xml", "g1.xml"))
    if not os.path.exists(xml_path):
        _G1_JOINT_AXIS_INDEX_CACHE = {}
        return _G1_JOINT_AXIS_INDEX_CACHE

    tree = ET.parse(xml_path)
    root = tree.getroot()

    joint_axes = {}
    for xml_class in tree.findall(".//default"):
        if "class" not in xml_class.attrib:
            continue
        joint_nodes = xml_class.findall("joint")
        if joint_nodes:
            joint_axes[xml_class.get("class")] = joint_nodes[0].get("axis")

    # mujoco (z-up, x-forward) -> ardy (y-up, z-forward)
    mujoco_to_ardy = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    axis_indices_by_name: dict[str, int] = {}
    for joint in root.find("worldbody").findall(".//joint"):
        axis_str = joint.get("axis") or joint_axes.get(joint.get("class"))
        if axis_str is None:
            continue
        axis_vals = np.array([float(x) for x in axis_str.split()], dtype=np.float64)
        if not np.any(axis_vals):
            continue
        axis_ardy = mujoco_to_ardy @ axis_vals
        axis_idx = int(np.argmax(np.abs(axis_ardy)))
        axis_indices_by_name[joint.get("name").replace("_joint", "_skel")] = axis_idx

    _G1_JOINT_AXIS_INDEX_CACHE = axis_indices_by_name
    return _G1_JOINT_AXIS_INDEX_CACHE


def _skew(v: np.ndarray) -> np.ndarray:
    """Skew-symmetric matrix for cross products: skew(v) @ x == np.cross(v, x)."""
    vx, vy, vz = float(v[0]), float(v[1]), float(v[2])
    return np.array([[0.0, -vz, vy], [vz, 0.0, -vx], [-vy, vx, 0.0]], dtype=np.float64)


def _rotation_matrix_from_two_vec(v_from: np.ndarray, v_to: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    """Return R such that R @ v_from ~= v_to (both treated as 3D vectors).

    Uses a Rodrigues-style construction, with special handling for near-parallel and near-opposite
    vectors for numerical stability.
    """
    a = np.asarray(v_from, dtype=np.float64).reshape(3)
    b = np.asarray(v_to, dtype=np.float64).reshape(3)
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na < eps or nb < eps:
        return np.eye(3, dtype=np.float64)
    a = a / na
    b = b / nb

    c = float(np.clip(np.dot(a, b), -1.0, 1.0))  # cos(theta)
    if c > 1.0 - eps:
        return np.eye(3, dtype=np.float64)
    if c < -1.0 + eps:
        # 180 deg rotation about any axis orthogonal to a:
        # R = -I + 2 * uu^T, where u is a unit axis orthogonal to a.
        axis_seed = np.array([1.0, 0.0, 0.0], dtype=np.float64)
        if abs(float(np.dot(a, axis_seed))) > 0.9:
            axis_seed = np.array([0.0, 1.0, 0.0], dtype=np.float64)
        u = np.cross(a, axis_seed)
        u = u / np.linalg.norm(u).clip(min=eps)
        return -np.eye(3, dtype=np.float64) + 2.0 * np.outer(u, u)

    v = np.cross(a, b)
    s2 = float(np.dot(v, v))  # ||v||^2 == sin^2(theta)
    K = _skew(v)
    # R = I + K + K^2 * ((1 - c) / s^2)
    return np.eye(3, dtype=np.float64) + K + (K @ K) * ((1.0 - c) / s2)


# TODO: should things in here by kept on cpu with numpy to avoid latency when interacting with UI?
#       the downside is we need torch/GPU for FK

# Cache arrow meshes to avoid recreating them repeatedly
_CACHED_ARROW_BASE = None
_CACHED_ARROW_HEAD = None


def _get_cached_arrow_meshes():
    """Get cached arrow base and head meshes (vertices and faces only)."""
    global _CACHED_ARROW_BASE, _CACHED_ARROW_HEAD
    if _CACHED_ARROW_BASE is None:
        arrow_base_mesh = trimesh.creation.cylinder(radius=0.01, height=0.2)
        arrow_head_mesh = trimesh.creation.cone(radius=0.03, height=0.05)
        _CACHED_ARROW_BASE = {
            "vertices": arrow_base_mesh.vertices.copy(),
            "faces": arrow_base_mesh.faces.copy(),
        }
        _CACHED_ARROW_HEAD = {
            "vertices": arrow_head_mesh.vertices.copy(),
            "faces": arrow_head_mesh.faces.copy(),
        }
    return _CACHED_ARROW_BASE, _CACHED_ARROW_HEAD


class WaypointMesh:
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        position: np.ndarray,
        heading: Optional[np.ndarray] = None,
        color: Optional[Tuple[int, int, int]] = (255, 0, 0),
        add_annulus: bool = True,
    ):
        self.server = server
        self.color = color
        self.base_name = name  # Store base name to remove parent folder later

        sphere = trimesh.creation.icosphere(subdivisions=3, radius=0.025)

        z_to_y_up = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]])

        self.sphere = self.server.scene.add_mesh_simple(
            name=f"{name}/sphere",
            vertices=sphere.vertices,
            faces=sphere.faces,
            position=position,
            color=color,
        )

        if add_annulus:
            annulus = trimesh.creation.annulus(r_min=0.1, r_max=0.2, height=0.005)
            annulus_vertices = annulus.vertices @ z_to_y_up
            self.annulus = self.server.scene.add_mesh_simple(
                name=f"{name}/annulus",
                vertices=annulus_vertices,
                faces=annulus.faces,
                position=position,
                color=color,
            )
        else:
            self.annulus = None

        self.arrow_base = None
        self.arrow_head = None
        if heading is not None:
            assert heading.shape == (2,), "Heading must be a 2D vector"
            heading_norm = heading / np.linalg.norm(heading)
            heading_scaled = 0.2 * heading_norm
            heading_3d = np.array([heading_scaled[0], 0, heading_scaled[1]])

            # Calculate rotation to align Y-axis (default cylinder/cone orientation) with heading
            # Rotation angle around Y-axis
            angle = np.arctan2(heading_norm[0], heading_norm[1])  # heading = [cos, sin] -> angle
            cos_a, sin_a = np.cos(angle), np.sin(angle)
            # Rotation matrix around Y-axis
            rot_y = np.array([[cos_a, 0, sin_a], [0, 1, 0], [-sin_a, 0, cos_a]])

            # Use cached arrow meshes
            arrow_base_cache, arrow_head_cache = _get_cached_arrow_meshes()

            # Rotate vertices to align with heading direction
            arrow_base_vertices = arrow_base_cache["vertices"] @ rot_y.T
            arrow_head_vertices = arrow_head_cache["vertices"] @ rot_y.T

            self.arrow_base = self.server.scene.add_mesh_simple(
                name=f"{name}/arrow_base",
                vertices=arrow_base_vertices,
                faces=arrow_base_cache["faces"],
                position=position + (heading_3d / 2),
                color=color,
            )
            self.arrow_head = self.server.scene.add_mesh_simple(
                name=f"{name}/arrow_head",
                vertices=arrow_head_vertices,
                faces=arrow_head_cache["faces"],
                position=position + heading_3d,
                color=color,
            )

    def update_position(self, position: np.ndarray, heading: Optional[np.ndarray] = None):
        self.sphere.position = position
        if self.annulus is not None:
            self.annulus.position = position
        if heading is not None:
            assert heading.shape == (2,), "Heading must be a 2D vector"
            heading_norm = heading / np.linalg.norm(heading)
            heading_scaled = 0.2 * heading_norm
            heading_3d = np.array([heading_scaled[0], 0, heading_scaled[1]])

            # Calculate rotation to align Y-axis with heading
            angle = np.arctan2(heading_norm[0], heading_norm[1])
            cos_a, sin_a = np.cos(angle), np.sin(angle)
            rot_y = np.array([[cos_a, 0, sin_a], [0, 1, 0], [-sin_a, 0, cos_a]])

            # Remove old arrow meshes if they exist
            if self.arrow_base is not None:
                self.server.scene.remove_by_name(self.arrow_base.name)
            if self.arrow_head is not None:
                self.server.scene.remove_by_name(self.arrow_head.name)

            # Use cached arrow meshes
            arrow_base_cache, arrow_head_cache = _get_cached_arrow_meshes()

            # Rotate vertices to align with heading direction
            arrow_base_vertices = arrow_base_cache["vertices"] @ rot_y.T
            arrow_head_vertices = arrow_head_cache["vertices"] @ rot_y.T

            # Get base name from sphere (e.g., "/waypoint_0/sphere" -> "/waypoint_0")
            base_name = self.sphere.name.rsplit("/", 1)[0]

            self.arrow_base = self.server.scene.add_mesh_simple(
                name=f"{base_name}/arrow_base",
                vertices=arrow_base_vertices,
                faces=arrow_base_cache["faces"],
                position=position + (heading_3d / 2),
                color=self.color,
            )
            self.arrow_head = self.server.scene.add_mesh_simple(
                name=f"{base_name}/arrow_head",
                vertices=arrow_head_vertices,
                faces=arrow_head_cache["faces"],
                position=position + heading_3d,
                color=self.color,
            )

    def clear(self):
        """Remove all waypoint meshes and the parent folder from the scene."""
        # self.server.scene.remove_by_name(self.sphere.name)
        # if self.annulus is not None:
        #     self.server.scene.remove_by_name(self.annulus.name)
        # if self.arrow_base is not None:
        #     self.server.scene.remove_by_name(self.arrow_base.name)
        # if self.arrow_head is not None:
        #     self.server.scene.remove_by_name(self.arrow_head.name)
        # Remove the parent folder to prevent empty folders in scene tree
        try:
            # print(f"Removing parent folder {self.base_name}")
            self.server.scene.remove_by_name(self.base_name)
        except Exception as e:
            print(f"Error removing parent folder {self.base_name}: {e}")
            pass  # Parent folder might have already been removed


class VelocityArrowMesh:
    """Visualizes root velocity as an arrow (line segment + cone head)."""

    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        color: tuple = (50, 150, 255),  # Default blue color (RGB)
    ):
        """Initialize velocity arrow visualization.

        Args:
            name: str, base name for the arrow components
            server: viser.ViserServer, server to add the arrow to
            skeleton: SkeletonBase, skeleton to get root index from
            color: tuple, RGB color tuple (0-255), default is blue
        """
        self.name = name
        self.server = server
        self.skeleton = skeleton
        self.color = color

        # Arrow components
        self.arrow_line = None  # Line segment
        self.arrow_cone = None  # Cone head
        self.should_show = False  # Track if arrow should be visible based on velocity magnitude

    def update(
        self,
        root_velocity: Optional[Union[np.ndarray, torch.Tensor]],
        root_pos: Union[np.ndarray, torch.Tensor],
        visible: bool = True,
    ):
        """Update the velocity arrow visualization.

        Args:
            root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s
            root_pos: Union[np.ndarray, torch.Tensor], [3] root position
            visible: bool, whether the arrow should be visible (controlled by skeleton visibility)
        """
        if root_velocity is None:
            # Hide arrow if no velocity provided
            self.should_show = False
            if self.arrow_line is not None:
                self.arrow_line.visible = False
                self.arrow_cone.visible = False
            return

        # Convert to numpy if tensor
        if isinstance(root_velocity, torch.Tensor):
            root_velocity = root_velocity.detach().cpu().numpy()
        if isinstance(root_pos, torch.Tensor):
            root_pos = root_pos.detach().cpu().numpy()

        # Project velocity to XZ plane
        velocity_xz = np.array([root_velocity[0], 0.0, root_velocity[2]])
        velocity_magnitude = np.linalg.norm(velocity_xz)

        # Only show arrow if velocity is significant (> 0.1 m/s)
        if velocity_magnitude <= 0.1:
            self.should_show = False
            if self.arrow_line is not None:
                self.arrow_line.visible = False
                self.arrow_cone.visible = False
            return

        # Arrow should be shown
        self.should_show = True

        # Calculate arrow geometry
        root_2d_pos = np.array([root_pos[0], 0.0, root_pos[2]])  # Project to ground
        velocity_dir = velocity_xz / velocity_magnitude
        arrow_length = velocity_magnitude / 4.0
        arrow_end = root_2d_pos + velocity_dir * arrow_length

        # Calculate rotation quaternion for arrow head
        from_vec = np.array([0.0, 0.0, 1.0])
        to_vec = velocity_dir
        rot_mat = _rotation_matrix_from_two_vec(from_vec, to_vec)
        quat = tf.SO3.from_matrix(rot_mat).wxyz

        # Create or update arrow components
        if self.arrow_line is None:
            # Create line segment
            self.arrow_line = self.server.scene.add_line_segments(
                name=f"{self.name}/velocity_line",
                points=np.array([[root_2d_pos, arrow_end]]),
                colors=self.color,
                line_width=3.0,
            )

            # Create cone head
            arrow_head = trimesh.creation.cone(radius=0.04, height=0.1)
            self.arrow_cone = self.server.scene.add_mesh_simple(
                name=f"{self.name}/velocity_cone",
                vertices=arrow_head.vertices,
                faces=arrow_head.faces,
                color=self.color,
                position=arrow_end,
                wxyz=quat,
            )
        else:
            # Update existing components atomically
            # Hide both during update
            old_line_visible = self.arrow_line.visible
            old_cone_visible = self.arrow_cone.visible
            self.arrow_line.visible = False
            self.arrow_cone.visible = False

            # Update geometry
            self.arrow_line.points = np.array([[root_2d_pos, arrow_end]])
            self.arrow_cone.position = arrow_end - 0.05 * velocity_dir
            self.arrow_cone.wxyz = quat

            # Restore visibility
            self.arrow_line.visible = old_line_visible
            self.arrow_cone.visible = old_cone_visible

        # Set visibility based on skeleton visibility and should_show
        self.arrow_line.visible = visible and self.should_show
        self.arrow_cone.visible = visible and self.should_show

    def set_visibility(self, visible: bool):
        """Set visibility of the velocity arrow."""
        if self.arrow_line is not None:
            self.arrow_line.visible = visible and self.should_show
        if self.arrow_cone is not None:
            self.arrow_cone.visible = visible and self.should_show

    def clear(self):
        """Remove the velocity arrow from the scene."""
        if self.arrow_line is not None:
            self.server.scene.remove_by_name(self.arrow_line.name)
            self.arrow_line = None
        if self.arrow_cone is not None:
            self.server.scene.remove_by_name(self.arrow_cone.name)
            self.arrow_cone = None


class SkeletonMesh:
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        joint_color: Optional[Tuple[float, float, float] | np.ndarray] = (
            255,
            235,
            0,
        ),
        bone_color: Optional[Tuple[float, float, float] | np.ndarray] = (
            27,
            106,
            0,
        ),
        starting_joints_pos: Optional[torch.Tensor] = None,
        show_root_2d_projection: bool = False,
    ):
        """
        name: str, name of the skeleton mesh
        server: viser.ViserServer, server to add the skeleton mesh to
        skeleton: SkeletonBase, skeleton to visualize (must be a CoreSkeleton27)
        joint_color: Optional[Tuple[float, float, float] | np.ndarray], color of the joints, either (3,) or (J, 3)
        bone_color: Optional[Tuple[float, float, float] | np.ndarray], color of the bones, either (3,) or (J-1, 3)
        starting_joints_pos: Optional[torch.Tensor], starting joint positions (if None, will use neutral pose)
        show_root_2d_projection: bool, whether to show the 2D root projection as a blue sphere
        """
        self.server = server
        self.skeleton = skeleton
        self.show_root_2d_projection = show_root_2d_projection
        joint_mesh = trimesh.creation.icosphere(subdivisions=3, radius=0.02)
        bone_mesh = trimesh.creation.cylinder(radius=0.01, height=1.0)

        init_joints_pos = skeleton.neutral_joints.clone()
        self.num_joints = init_joints_pos.shape[0]
        num_bones = self.num_joints - 1
        non_root_bones = [
            joint_name
            for joint_name, parent_name in self.skeleton.bone_order_names_with_parents
            if parent_name is not None
        ]
        self.bone_to_idx = {bone_name: idx for idx, bone_name in enumerate(non_root_bones)}

        # initialize meshes
        init_joints_wxyzs = np.concatenate([np.ones((self.num_joints, 1)), np.zeros((self.num_joints, 3))], axis=1)
        if isinstance(joint_color, tuple):
            self.joint_colors = np.full((self.num_joints, 3), joint_color)
        elif isinstance(joint_color, np.ndarray):
            assert joint_color.shape == (
                self.num_joints,
                3,
            ), "Joint colors must be (J, 3)"
            self.joint_colors = joint_color
        self.joints_batched_mesh = server.scene.add_batched_meshes_simple(
            f"{name}/joints",
            vertices=joint_mesh.vertices,
            faces=joint_mesh.faces,
            batched_wxyzs=init_joints_wxyzs,
            batched_positions=np.zeros((self.num_joints, 3)),
            batched_scales=np.ones((self.num_joints, 3)),
            batched_colors=self.joint_colors,
        )
        init_bones_wxyzs = np.concatenate([np.ones((num_bones, 1)), np.zeros((num_bones, 3))], axis=1)
        if isinstance(bone_color, tuple):
            bone_color = np.full((num_bones, 3), bone_color)
        elif isinstance(bone_color, np.ndarray):
            assert bone_color.shape == (num_bones, 3), "Bone colors must be (J-1, 3)"
            bone_color = bone_color
        self.bones_batched_mesh = server.scene.add_batched_meshes_simple(
            f"{name}/bones",
            vertices=bone_mesh.vertices,
            faces=bone_mesh.faces,
            batched_wxyzs=init_bones_wxyzs,
            batched_positions=np.zeros((num_bones, 3)),
            batched_scales=np.ones((num_bones, 3)),
            batched_colors=bone_color,
        )

        # Initialize 2D root projection sphere (blue)
        if self.show_root_2d_projection:
            root_2d_mesh = trimesh.creation.icosphere(subdivisions=2, radius=0.03)
            self.root_2d_sphere = server.scene.add_mesh_simple(
                f"{name}/root_2d_projection",
                vertices=root_2d_mesh.vertices,
                faces=root_2d_mesh.faces,
                color=(50, 150, 255),  # Blue color
                position=(0, 0, 0),
            )
        else:
            self.root_2d_sphere = None

        # Initialize root velocity arrow (visualize XZ projection)
        self.velocity_arrow_mesh = VelocityArrowMesh(
            name=name,
            server=server,
            skeleton=skeleton,
        )

        # used if precomputed meshes are used
        self.mesh_info_cache = None

        if starting_joints_pos is not None:
            self.set_pose(starting_joints_pos)
        else:
            # set them to neutral pose
            min_height = init_joints_pos[:, 1].min().item()
            init_joints_pos[:, 1] -= min_height  # move to be on ground
            self.set_pose(init_joints_pos)

    def compute_single_pose(self, joints_pos: np.ndarray):
        """Compute the mesh for a single frame.

        joints_pos: [J, 3] global joint positions
        """
        # compute bone transforms
        new_batched_positions = np.zeros((self.skeleton.nbjoints - 1, 3))
        new_batched_wxyzs = np.zeros((self.skeleton.nbjoints - 1, 4))
        new_batched_scales = np.ones((self.skeleton.nbjoints - 1, 3))
        for joint_name, parent_name in self.skeleton.bone_order_names_with_parents:
            if parent_name is None:
                continue
            joint_idx = self.skeleton.bone_index[joint_name]
            parent_idx = self.skeleton.bone_index[parent_name]
            joint_pos = joints_pos[joint_idx]
            parent_pos = joints_pos[parent_idx]

            bone_pos = (joint_pos + parent_pos) / 2.0
            bone_scale = np.linalg.norm(joint_pos - parent_pos)
            if bone_scale < 1e-8:
                bone_wxyz = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
            else:
                bone_dir = (joint_pos - parent_pos) / bone_scale
                R = _rotation_matrix_from_two_vec(np.array([0.0, 0.0, 1.0], dtype=np.float64), bone_dir)
                bone_wxyz = tf.SO3.from_matrix(R).wxyz

            bone_idx = self.bone_to_idx[joint_name]
            new_batched_positions[bone_idx] = bone_pos
            new_batched_wxyzs[bone_idx] = bone_wxyz
            new_batched_scales[bone_idx] = np.array([1.0, 1.0, bone_scale], dtype=float)

        return new_batched_positions, new_batched_wxyzs, new_batched_scales

    def precompute_mesh_info(self, joints_pos: torch.Tensor):
        """Precompute the meshes for all frames at once.

        joints_pos: [T, J, 3] global joint positions
        """
        joints_pos = joints_pos.cpu().numpy()
        # compute bone transforms
        num_frames = joints_pos.shape[0]
        self.mesh_info_cache = {
            "positions": np.zeros((num_frames, self.skeleton.nbjoints - 1, 3)),
            "wxyzs": np.zeros((num_frames, self.skeleton.nbjoints - 1, 4)),
            "scales": np.ones((num_frames, self.skeleton.nbjoints - 1, 3)),
        }
        for i in range(num_frames):
            new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose(joints_pos[i])
            self.mesh_info_cache["positions"][i] = new_batched_positions
            self.mesh_info_cache["wxyzs"][i] = new_batched_wxyzs
            self.mesh_info_cache["scales"][i] = new_batched_scales

    def update_mesh_info_cache(self, joints_pos: torch.Tensor, frame_idx: int):
        """Update the mesh info cache for the given frame.

        joints_pos: [J, 3] global joint positions
        """
        assert self.mesh_info_cache is not None
        new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose(
            joints_pos.cpu().numpy()
        )
        self.mesh_info_cache["positions"][frame_idx] = new_batched_positions
        self.mesh_info_cache["wxyzs"][frame_idx] = new_batched_wxyzs
        self.mesh_info_cache["scales"][frame_idx] = new_batched_scales

    def set_pose(
        self,
        joints_pos: torch.Tensor,
        foot_contacts: Optional[torch.Tensor] = None,
        frame_idx: Optional[int] = None,
        root_velocity: Optional[Union[np.ndarray, torch.Tensor]] = None,
    ):
        """
        joints_pos: [J, 3] global joint positions
        foot_contacts: [4] contact labels for left heel/toe and right heel/toe, 1 for in contact
        frame_idx: int, index of the frame to set the pose for (only needed if precomputed mesh info is used)
        root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s
        """
        self.cur_joints_pos = joints_pos
        joints_pos = joints_pos.cpu().numpy()

        if self.mesh_info_cache is not None:
            assert frame_idx is not None
            new_batched_positions = self.mesh_info_cache["positions"][frame_idx]
            new_batched_wxyzs = self.mesh_info_cache["wxyzs"][frame_idx]
            new_batched_scales = self.mesh_info_cache["scales"][frame_idx]
        else:
            new_batched_positions, new_batched_wxyzs, new_batched_scales = self.compute_single_pose(joints_pos)

        # update meshes
        self.bones_batched_mesh.batched_positions = new_batched_positions
        self.bones_batched_mesh.batched_wxyzs = new_batched_wxyzs
        self.bones_batched_mesh.batched_scales = new_batched_scales
        # directly set joint positions
        self.joints_batched_mesh.batched_positions = joints_pos

        # update 2D root projection sphere
        if self.root_2d_sphere is not None:
            root_pos = joints_pos[self.skeleton.root_idx]
            root_2d_pos = np.array([root_pos[0], 0.0, root_pos[2]])
            self.root_2d_sphere.position = root_2d_pos

        # Update root velocity arrow visualization (XZ projection)
        root_pos = joints_pos[self.skeleton.root_idx]
        skeleton_visible = self.joints_batched_mesh.visible
        self.velocity_arrow_mesh.update(
            root_velocity=root_velocity,
            root_pos=root_pos,
            visible=skeleton_visible,
        )

        # update colors for foot contacts
        if foot_contacts is not None:
            cur_joint_colors = self.joint_colors.copy()
            foot_contacts = foot_contacts.bool().cpu().numpy().astype(bool)
            foot_joints = np.array(self.skeleton.foot_joint_idx, dtype=int)
            contact_idx = foot_joints[foot_contacts]
            cur_joint_colors[contact_idx] = (160, 32, 240)
            self.joints_batched_mesh.batched_colors = cur_joint_colors
        else:
            self.joints_batched_mesh.batched_colors = self.joint_colors

    def set_visibility(self, visible: bool):
        self.joints_batched_mesh.visible = visible
        self.bones_batched_mesh.visible = visible
        if self.root_2d_sphere is not None:
            self.root_2d_sphere.visible = visible
        # Update velocity arrow visibility
        self.velocity_arrow_mesh.set_visibility(visible)

    def get_pose(self) -> np.ndarray:
        return self.cur_joints_pos

    def clear(self):
        names = [mesh.name for mesh in [self.joints_batched_mesh, self.bones_batched_mesh]]
        for name in names:
            self.server.scene.remove_by_name(name)
        if self.root_2d_sphere is not None:
            self.server.scene.remove_by_name(self.root_2d_sphere.name)
        # Clear velocity arrow
        self.velocity_arrow_mesh.clear()


LIGHT_THEME = dict(
    mesh=(152, 189, 255),  # (90, 200, 255) - original viser blue
)

DARK_THEME = dict(
    mesh=(60, 85, 130),
)


class Character:
    def __init__(
        self,
        name: str,
        server: viser.ViserServer | viser.ClientHandle,
        skeleton: SkeletonBase,
        create_skeleton_mesh: bool = True,
        create_skinned_mesh: bool = True,
        visible_skeleton: bool = False,
        visible_skinned_mesh: bool = True,
        skinned_mesh_opacity: float = 1.0,
        show_foot_contacts: bool = True,
        dark_mode: bool = False,
        mesh_mode: str = "core_skin",
        g1_mesh_dir: Optional[str] = None,
        show_root_2d_projection: bool = False,
    ):
        self.server = server
        self.name = name
        self.skeleton = skeleton
        self.cur_joints_pos = None
        self.cur_joints_rot = None
        self.cur_foot_contacts = None

        self.skeleton_mesh = None
        self.show_foot_contacts = show_foot_contacts
        if create_skeleton_mesh:
            self.skeleton_mesh = SkeletonMesh(
                f"/{name}/skeleton",
                server,
                skeleton,
                show_root_2d_projection=show_root_2d_projection,
            )
            # init with default rest pose
            self.cur_joints_pos = self.skeleton_mesh.get_pose()
            self.skeleton_mesh.set_visibility(visible_skeleton)

        self.skinned_mesh = None
        self.g1_mesh_rig = None
        self.skin = None
        self.mesh_mode = mesh_mode
        if create_skinned_mesh:
            if isinstance(self.skeleton, CoreSkeleton27) and mesh_mode == "core_skin":
                self.skin = CoreSkin(self.skeleton)
                self.skinned_mesh = server.scene.add_mesh_simple(
                    f"/{name}/simple_skinned",
                    vertices=self.skin.bind_vertices.cpu().numpy(),
                    faces=self.skin.faces.cpu().numpy(),
                    opacity=None,
                    color=LIGHT_THEME["mesh"] if not dark_mode else DARK_THEME["mesh"],
                    wireframe=False,
                    visible=False,
                )
                self.skinned_verts_cache = None

                bind_pos = self.skeleton.neutral_joints.clone()
                min_height = bind_pos[:, 1].min().item()
                bind_pos[:, 1] -= min_height  # move to be on ground
                bind_rotmat = torch.eye(3, device=bind_pos.device).repeat(bind_pos.shape[0], 1, 1)
                self.set_pose(bind_pos, bind_rotmat)
                self.skinned_mesh.visible = True  # avoid blinking
                self.set_skinned_mesh_visibility(visible_skinned_mesh)
                self.set_skinned_mesh_opacity(skinned_mesh_opacity)
            elif isinstance(self.skeleton, (SOMASkeleton30, SOMASkeleton77)) and mesh_mode == "soma_skin":
                self.skin = SOMASkin(self.skeleton)
                self.skinned_mesh = server.scene.add_mesh_simple(
                    f"/{name}/simple_skinned",
                    vertices=self.skin.bind_vertices.cpu().numpy(),
                    faces=self.skin.faces.cpu().numpy(),
                    opacity=None,
                    color=LIGHT_THEME["mesh"] if not dark_mode else DARK_THEME["mesh"],
                    wireframe=False,
                    visible=False,
                )
                self.skinned_verts_cache = None

                bind_pos = self.skeleton.neutral_joints.clone()
                min_height = bind_pos[:, 1].min().item()
                bind_pos[:, 1] -= min_height  # move to be on ground
                bind_rotmat = torch.eye(3, device=bind_pos.device).repeat(bind_pos.shape[0], 1, 1)
                self.set_pose(bind_pos, bind_rotmat)
                self.skinned_mesh.visible = True  # avoid blinking
                self.set_skinned_mesh_visibility(visible_skinned_mesh)
                self.set_skinned_mesh_opacity(skinned_mesh_opacity)
            elif isinstance(self.skeleton, G1Skeleton34) and mesh_mode == "g1_stl":
                if g1_mesh_dir is None:
                    g1_mesh_dir = os.path.join(
                        os.path.dirname(__file__),
                        "..",
                        "assets",
                        "skeletons",
                        "g1skel34",
                        "meshes",
                        "g1",
                    )
                    g1_mesh_dir = os.path.abspath(g1_mesh_dir)
                if not os.path.exists(g1_mesh_dir):
                    print(f"G1 mesh directory not found: {g1_mesh_dir}")
                self.g1_mesh_rig = G1MeshRig(
                    name,
                    server,
                    self.skeleton,
                    g1_mesh_dir,
                    DARK_THEME["mesh"] if dark_mode else LIGHT_THEME["mesh"],
                )
                init_joints_pos = self.skeleton.neutral_joints.clone()
                min_height = init_joints_pos[:, 1].min().item()
                init_joints_pos[:, 1] -= min_height  # move to be on ground
                init_joints_rot = torch.eye(3, device=init_joints_pos.device).repeat(init_joints_pos.shape[0], 1, 1)
                self.set_pose(init_joints_pos, init_joints_rot)
                self.set_skinned_mesh_visibility(visible_skinned_mesh)
                self.set_skinned_mesh_opacity(skinned_mesh_opacity)
            else:
                raise ValueError(
                    "Unsupported mesh mode for skeleton type: "
                    f"{type(self.skeleton).__name__} with mesh_mode={mesh_mode}"
                )

    def change_theme(self, is_dark_mode):
        color = DARK_THEME["mesh"] if is_dark_mode else LIGHT_THEME["mesh"]
        if self.skinned_mesh is not None:
            self.skinned_mesh.color = color
        if self.g1_mesh_rig is not None:
            self.g1_mesh_rig.set_color(color)

    def set_skeleton_visibility(self, visible: bool):
        if self.skeleton_mesh is not None:
            self.skeleton_mesh.set_visibility(visible)

    def set_show_foot_contacts(self, show: bool):
        self.show_foot_contacts = show

    def set_skinned_mesh_visibility(self, visible: bool):
        if self.skinned_mesh is not None:
            self.skinned_mesh.visible = visible
        if self.g1_mesh_rig is not None:
            self.g1_mesh_rig.set_visibility(visible)

    def set_skinned_mesh_opacity(self, opacity: float):
        if self.skinned_mesh is not None:
            self.skinned_mesh.opacity = opacity
        if self.g1_mesh_rig is not None:
            self.g1_mesh_rig.set_opacity(opacity)

    def set_skinned_mesh_wireframe(self, wireframe: bool):
        if self.skinned_mesh is not None:
            self.skinned_mesh.wireframe = wireframe
        if self.g1_mesh_rig is not None:
            self.g1_mesh_rig.set_wireframe(wireframe)

    def precompute_skinning(self, joints_pos: torch.Tensor, joints_rot: torch.Tensor):
        """If using simple skinning, we can precompute the skinning for all frames at once.

        joints_pos: [T, J, 3] global joint positions
        joints_rot: [T, J, 3, 3] global joint rotation matrices
        """
        assert self.skin is not None
        self.skinned_verts_cache = self.skin.skin(joints_rot, joints_pos, rot_is_global=True).cpu().numpy()

    def update_skinning_cache(self, joints_pos: torch.Tensor, joints_rot: torch.Tensor, frame_idx: int):
        """Update the skinning cache for the given frame.

        joints_pos: [J, 3] global joint positions
        joints_rot: [J, 3, 3] global joint rotation matrices
        frame_idx: int, index of the frame to update the cache for
        """
        if self.skinned_verts_cache is None:
            return

        new_skinned_verts = self.skin.skin(joints_rot[None], joints_pos[None], rot_is_global=True)[0].cpu().numpy()
        self.skinned_verts_cache[frame_idx] = new_skinned_verts

    def set_pose(
        self,
        joints_pos: torch.Tensor,
        joints_rot: torch.Tensor,
        foot_contacts: Optional[torch.Tensor] = None,
        frame_idx: Optional[int] = None,
        root_velocity: Optional[Union[np.ndarray, torch.Tensor]] = None,
    ):
        """
        joints_pos: [J, 3] global joint positions
        joints_rot: [J, 3, 3] global joint rotation matrices
        frame_idx: int, index of the frame to set the pose for (only needed if precomputed skinning is used)
        root_velocity: Optional[Union[np.ndarray, torch.Tensor]], [3] root joint velocity (x, y, z) in m/s
        """
        if self.skeleton_mesh is not None:
            cur_foot_contacts = foot_contacts if self.show_foot_contacts else None
            self.skeleton_mesh.set_pose(
                joints_pos,
                foot_contacts=cur_foot_contacts,
                frame_idx=frame_idx,
                root_velocity=root_velocity,
            )
            self.cur_foot_contacts = cur_foot_contacts

        if self.skinned_mesh is not None:
            if self.skinned_verts_cache is not None:
                assert frame_idx is not None
                skinned_verts = self.skinned_verts_cache[frame_idx]
            else:
                skinned_verts = self.skin.skin(joints_rot[None], joints_pos[None], rot_is_global=True)[0].cpu().numpy()

            # update the vertices
            self.skinned_mesh.vertices = skinned_verts
        if self.g1_mesh_rig is not None:
            joints_pos_np = joints_pos.detach().cpu().numpy()
            joints_rot_np = joints_rot.detach().cpu().numpy()
            self.g1_mesh_rig.set_pose(joints_pos_np, joints_rot_np)

        self.cur_joints_pos = joints_pos
        self.cur_joints_rot = joints_rot

    def get_pose(self) -> torch.Tensor:
        return self.cur_joints_pos, self.cur_joints_rot

    def clear(self):
        if self.skeleton_mesh is not None:
            self.skeleton_mesh.clear()
        if self.skinned_mesh is not None:
            self.server.scene.remove_by_name(self.skinned_mesh.name)
        if self.g1_mesh_rig is not None:
            self.g1_mesh_rig.clear()


class CharacterMotion:
    def __init__(
        self,
        character: Character,
        joints_pos: torch.Tensor,
        joints_rot: torch.Tensor,
        foot_contacts: Optional[torch.Tensor] = None,
    ):
        self.character = character
        self.server = character.server
        self.skeleton = character.skeleton
        self.name = character.name

        # [T, J, 3] global joint positions
        self.joints_pos = joints_pos
        # [T, J, 3, 3] global joint rotation matrices
        self.joints_rot = joints_rot
        assert joints_pos.shape[0] == joints_rot.shape[0]
        # keep track of local rots as well for convenience during pose editing
        self.joints_local_rot = global_rots_to_local_rots(joints_rot, self.skeleton)

        self.length = joints_pos.shape[0]
        self.cur_frame_idx = None

        self.foot_contacts = foot_contacts
        if foot_contacts is not None:
            assert foot_contacts.shape[0] == self.length

        self.precompute_mesh_info()

        # gizmos for pose editing
        self.root_translation_gizmo = None
        self.updating_root_translation_gizmo = False
        self.joint_gizmos = None
        self.updating_joint_gizmos = False

    def precompute_mesh_info(self):
        if self.character.skeleton_mesh is not None:
            print("Caching skeleton mesh info...")
            self.character.skeleton_mesh.precompute_mesh_info(self.joints_pos)
        if self.character.skinned_mesh is not None:
            print("Caching skinning info...")
            self.character.precompute_skinning(self.joints_pos, self.joints_rot)

    def set_frame(self, idx: int):
        """Sets the pose of the character to the given frame index."""
        idx = min(idx, self.length - 1)  # clamp to last frame
        cur_foot_contacts = self.foot_contacts[idx] if self.foot_contacts is not None else None
        self.character.set_pose(
            self.joints_pos[idx],
            self.joints_rot[idx],
            frame_idx=idx,
            foot_contacts=cur_foot_contacts,
        )
        self.cur_frame_idx = idx

        # update gizmos if frame has changed due to playback
        cur_root_pos = self.joints_pos[self.cur_frame_idx, self.skeleton.root_idx].clone()
        cur_root_pos[1] = 0.0
        if self.root_translation_gizmo is not None and not self.updating_root_translation_gizmo:
            self.root_translation_gizmo.position = cur_root_pos.cpu().numpy()
        if self.joint_gizmos is not None:
            for i, joint_gizmo in enumerate(self.joint_gizmos):
                if not self.updating_joint_gizmos:
                    joint_gizmo.position = self.joints_pos[self.cur_frame_idx, i].cpu().numpy()
                    joint_gizmo.wxyz = tf.SO3.from_matrix(
                        self.joints_local_rot[self.cur_frame_idx, i].cpu().numpy()
                    ).wxyz

    def update_pose_at_frame(
        self,
        frame_idx: int,
        joints_pos: Optional[torch.Tensor] = None,
        joints_rot: Optional[torch.Tensor] = None,
        joints_local_rot: Optional[torch.Tensor] = None,
        foot_contacts: Optional[torch.Tensor] = None,
    ):
        """Overwrites one or more of the pose components at the given frame.

        If only a subset of joints_pos, joints_rot, or joints_local_rot are provided, the other
        components will be updated with FK.
        """
        if joints_pos is not None:
            joints_pos = to_torch(joints_pos, device=self.joints_pos.device, dtype=self.joints_pos.dtype)
            self.joints_pos[frame_idx] = joints_pos
            if joints_local_rot is None and joints_rot is None:
                raise NotImplementedError("No IK to update joint rotations accordingly.")
        if joints_rot is not None:
            joints_rot = to_torch(joints_rot, device=self.joints_rot.device, dtype=self.joints_rot.dtype)
            self.joints_rot[frame_idx] = joints_rot
            if joints_local_rot is None:
                # update local rots from global rots
                self.joints_local_rot[frame_idx] = global_rots_to_local_rots(joints_rot, self.skeleton)
            if joints_pos is None:
                # need to update with FK
                new_posed_joints, _ = batch_rigid_transform(
                    self.joints_local_rot[frame_idx : frame_idx + 1],
                    self.skeleton.neutral_joints[None].to(self.joints_local_rot.device),
                    self.skeleton.joint_parents.to(self.joints_local_rot.device),
                    self.skeleton.root_idx,
                )
                new_posed_joints = (
                    new_posed_joints[0]
                    + self.joints_pos[frame_idx, self.skeleton.root_idx : self.skeleton.root_idx + 1]
                    - self.skeleton.neutral_joints[[self.skeleton.root_idx]]
                )
                self.joints_pos[frame_idx] = new_posed_joints
        if joints_local_rot is not None:
            joints_local_rot = to_torch(joints_local_rot, device=self.joints_local_rot.device).to(
                dtype=self.joints_local_rot.dtype
            )
            self.joints_local_rot[frame_idx] = joints_local_rot
            if joints_rot is None or joints_pos is None:
                # need to update with FK
                new_posed_joints, new_global_rots = batch_rigid_transform(
                    self.joints_local_rot[frame_idx : frame_idx + 1],
                    self.skeleton.neutral_joints[None].to(self.joints_local_rot.device),
                    self.skeleton.joint_parents.to(self.joints_local_rot.device),
                    self.skeleton.root_idx,
                )
                new_posed_joints = (
                    new_posed_joints[0]
                    + self.joints_pos[frame_idx, self.skeleton.root_idx : self.skeleton.root_idx + 1]
                    - self.skeleton.neutral_joints[[self.skeleton.root_idx]]
                )
                if joints_rot is None:
                    self.joints_rot[frame_idx] = new_global_rots[0]
                if joints_pos is None:
                    self.joints_pos[frame_idx] = new_posed_joints
        if foot_contacts is not None:
            foot_contacts = to_torch(foot_contacts, device=self.foot_contacts.device).to(dtype=self.foot_contacts.dtype)
            self.foot_contacts[frame_idx] = foot_contacts

        if self.character.skeleton_mesh is not None:
            self.character.skeleton_mesh.update_mesh_info_cache(self.joints_pos[frame_idx], frame_idx)
        if self.character.skinned_mesh is not None:
            self.character.update_skinning_cache(self.joints_pos[frame_idx], self.joints_rot[frame_idx], frame_idx)

    def clear(self):
        self.character.clear()

    #
    # Editing helpers
    #
    def get_current_projected_root_pos(self) -> np.ndarray:
        """Get the projected root position on the ground at the current frame."""
        root_pos = self.joints_pos[self.cur_frame_idx, self.skeleton.root_idx].clone()
        root_pos[1] = 0.0
        return to_numpy(root_pos)

    def get_projected_root_pos(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray:
        """If requested frames are out of range, simply pads with the last frame to get expected
        length."""
        if end_frame_idx is None:
            expected_len = 1
        else:
            expected_len = end_frame_idx - start_frame_idx + 1
        if start_frame_idx >= self.length:
            start_frame_idx = self.length - 1
        if end_frame_idx is None or expected_len == 1:
            root_pos = self.joints_pos[start_frame_idx, self.skeleton.root_idx].clone()
            root_pos[1] = 0.0
            return to_numpy(root_pos)
        else:
            if end_frame_idx >= self.length:
                end_frame_idx = self.length = 1
            root_pos = self.joints_pos[start_frame_idx : end_frame_idx + 1, self.skeleton.root_idx].clone()
            root_pos[:, 1] = 0.0
            if root_pos.shape[0] < expected_len:
                # pad with the last root position
                root_pos = torch.cat(
                    [
                        root_pos,
                        root_pos[-1:].repeat(expected_len - root_pos.shape[0], 1),
                    ],
                    dim=0,
                )
            return to_numpy(root_pos)

    def set_projected_root_pos_path(
        self,
        root_pos_path: np.ndarray | torch.Tensor,
        min_frame_idx: int = None,
        max_frame_idx: int = None,
    ):
        """Sets the projected root position path for the character motion.

        Can set only a subset of the path by providing min_frame_idx and max_frame_idx. If not provided, will set the full
        path.

        Args:
            root_pos_path: torch.Tensor, [T, 2] projected root positions
            min_frame_idx: int, optional, minimum frame index to set the path at
            max_frame_idx: int, optional, maximum frame index to set the path at
        """
        if min_frame_idx is not None or max_frame_idx is not None:
            assert min_frame_idx is not None and max_frame_idx is not None, (
                "min_frame_idx and max_frame_idx must be provided if setting path at specific frames"
            )
            if min_frame_idx >= self.length:
                # both are out of bounds
                return
            max_frame_idx = min(max_frame_idx, self.length - 1)
            root_pos_path = root_pos_path[min_frame_idx : max_frame_idx + 1]
        else:
            assert root_pos_path.shape[0] == self.length
            min_frame_idx = 0
            max_frame_idx = self.length - 1

        cur_joints_pos = self.joints_pos.clone()[min_frame_idx : max_frame_idx + 1]
        root_pos_tensor = to_torch(root_pos_path, device=cur_joints_pos.device, dtype=cur_joints_pos.dtype)
        diff = root_pos_tensor - cur_joints_pos[:, self.skeleton.root_idx, [0, 2]]
        cur_joints_pos[:, :, [0, 2]] += diff.unsqueeze(1)
        for frame_idx in range(min_frame_idx, max_frame_idx + 1):
            rel_idx = frame_idx - min_frame_idx
            self.update_pose_at_frame(
                frame_idx,
                joints_pos=cur_joints_pos[rel_idx],
                joints_rot=self.joints_rot[frame_idx],
                joints_local_rot=self.joints_local_rot[frame_idx],
            )
        # update immediately to show changes
        self.set_frame(self.cur_frame_idx)

    def get_joints_pos(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray:
        """If requested frames are out of range, simply pads with the last frame to get expected
        length."""
        if end_frame_idx is None:
            expected_len = 1
        else:
            expected_len = end_frame_idx - start_frame_idx + 1
        if start_frame_idx >= self.length:
            start_frame_idx = self.length - 1
        if end_frame_idx is None or expected_len == 1:
            return to_numpy(self.joints_pos[start_frame_idx].clone())
        else:
            if end_frame_idx >= self.length:
                end_frame_idx = self.length - 1
            return_joints_pos = self.joints_pos[start_frame_idx : end_frame_idx + 1].clone()
            if return_joints_pos.shape[0] < expected_len:
                # pad with the last pose
                return_joints_pos = torch.cat(
                    [
                        return_joints_pos,
                        return_joints_pos[-1:].repeat(expected_len - return_joints_pos.shape[0], 1, 1),
                    ],
                    dim=0,
                )
            return to_numpy(return_joints_pos)

    def get_joints_rot(self, start_frame_idx: int, end_frame_idx: int = None) -> np.ndarray:
        """If requested frames are out of range, simply pads with the last frame to get expected
        length."""
        if end_frame_idx is None:
            expected_len = 1
        else:
            expected_len = end_frame_idx - start_frame_idx + 1
        if start_frame_idx >= self.length:
            start_frame_idx = self.length - 1
        if end_frame_idx is None or expected_len == 1:
            return to_numpy(self.joints_rot[start_frame_idx].clone())
        else:
            if end_frame_idx >= self.length:
                end_frame_idx = self.length - 1
            return_joints_rot = self.joints_rot[start_frame_idx : end_frame_idx + 1].clone()
            if return_joints_rot.shape[0] < expected_len:
                # pad with the last pose
                return_joints_rot = torch.cat(
                    [
                        return_joints_rot,
                        return_joints_rot[-1:].repeat(expected_len - return_joints_rot.shape[0], 1, 1, 1),
                    ],
                    dim=0,
                )
            return to_numpy(return_joints_rot)

    def get_current_joints_pos(self) -> torch.Tensor:
        return self.joints_pos[self.cur_frame_idx].clone()

    def get_current_joints_rot(self) -> torch.Tensor:
        return self.joints_rot[self.cur_frame_idx].clone()

    def add_root_translation_gizmo(self, constraints: dict):
        """Create and initialize gizmo to control the root translation."""
        # TODO: could also allow rotation around y-axis
        self.root_translation_gizmo = self.server.scene.add_transform_controls(
            f"/{self.name}/gizmo_root_translation",
            scale=0.5,
            line_width=2.5,
            active_axes=(True, False, True),  # only allow translation on xz plane
            disable_axes=False,
            disable_sliders=False,
            disable_rotations=True,
            depth_test=False,  # render even when occluded
        )
        init_position = self.get_current_projected_root_pos()
        self.root_translation_gizmo.position = init_position

        @self.root_translation_gizmo.on_update
        def _(_):
            self.updating_root_translation_gizmo = True
            # translate to gizmo position
            new_root_pos = to_torch(
                self.root_translation_gizmo.position,
                device=self.joints_pos.device,
            ).to(dtype=self.joints_pos.dtype)
            cur_joints_pos = self.joints_pos[self.cur_frame_idx].clone()
            root_diff = new_root_pos - cur_joints_pos[self.skeleton.root_idx]
            root_diff[1] = 0.0  # don't change height
            cur_joints_pos += root_diff[None]
            self.update_pose_at_frame(
                self.cur_frame_idx,
                joints_pos=cur_joints_pos,
                joints_rot=self.joints_rot[self.cur_frame_idx],
                joints_local_rot=self.joints_local_rot[self.cur_frame_idx],
            )

            self.updating_root_translation_gizmo = False
            # update immediately to show user changes
            self.set_frame(self.cur_frame_idx)
            # update the 2D waypoint constraints as well if there is one
            if "2D Root" in constraints:
                root_2d_contraints = constraints["2D Root"]
                # if there is a constraint at that frame, we want to update it
                frame_idx = self.cur_frame_idx
                if frame_idx in root_2d_contraints.keyframes:
                    for keyframe_id in root_2d_contraints.frame2keyid[frame_idx]:
                        # add will modify the existing constraint
                        root_2d_contraints.add_keyframe(
                            keyframe_id,
                            frame_idx,
                            root_pos=new_root_pos,
                            exists_ok=True,
                        )
            if "Full-Body" in constraints:
                full_body_constraints = constraints["Full-Body"]
                # if there is a constraint at that frame, we want to update it
                frame_idx = self.cur_frame_idx
                if frame_idx in full_body_constraints.keyframes:
                    current_dict = full_body_constraints.keyframes[frame_idx]
                    for keyframe_id in full_body_constraints.frame2keyid[frame_idx]:
                        # add will modify the existing constraint
                        full_body_constraints.add_keyframe(
                            keyframe_id,
                            frame_idx,
                            joints_pos=cur_joints_pos,
                            joints_rot=current_dict["joints_rot"],
                            exists_ok=True,
                        )
            if "End-Effectors" in constraints:
                end_effector_constraints = constraints["End-Effectors"]
                # if there is a constraint at that frame, we want to update it
                frame_idx = self.cur_frame_idx
                if frame_idx in end_effector_constraints.keyframes:
                    current_dict = end_effector_constraints.keyframes[frame_idx]
                    for keyframe_id, _ in end_effector_constraints.frame2keyid[frame_idx]:
                        # add will modify the existing constraint
                        end_effector_constraints.add_keyframe(
                            keyframe_id,
                            frame_idx,
                            joints_pos=cur_joints_pos,
                            joints_rot=current_dict["joints_rot"],
                            joint_names=current_dict["joint_names"],
                            end_effector_type=current_dict["end_effector_type"],
                            exists_ok=True,
                        )

    def add_joint_gizmos(self, constraints: dict):
        self.joint_gizmos = []

        joint_axis_indices = None
        hidden_gizmo_joints = None
        if isinstance(self.skeleton, G1Skeleton34):
            joint_axis_indices = _get_g1_joint_axis_indices()
            hidden_gizmo_joints = set(
                self.skeleton.left_hand_joint_names
                + self.skeleton.right_hand_joint_names
                + self.skeleton.left_foot_joint_names
                + self.skeleton.right_foot_joint_names
            )
        elif isinstance(self.skeleton, CoreSkeleton27):
            hidden_gizmo_joints = {
                "RightHandThumb1",
                "RightHandEnd",
                "LeftHandThumb1",
                "LeftHandEnd",
            }

        joints_wxyzs = tf.SO3.from_matrix(self.joints_local_rot[self.cur_frame_idx].cpu().numpy()).wxyz
        for joint_idx in range(self.skeleton.nbjoints):
            disable_axes = True  # by default, only rotation controls
            disable_sliders = True
            if joint_idx == self.skeleton.root_idx:
                disable_axes = False  # allow translation for root
                disable_sliders = False
            active_axes = (True, True, True)
            if joint_axis_indices is not None:
                joint_name = self.skeleton.bone_order_names[joint_idx]
                axis_idx = joint_axis_indices.get(joint_name)
                if axis_idx is not None:
                    # PivotControls shows rotation handles when a plane is active.
                    # To allow rotation about one axis, enable the other two axes.
                    active_axes = (
                        axis_idx != 0,
                        axis_idx != 1,
                        axis_idx != 2,
                    )
            joint_visible = True
            if hidden_gizmo_joints is not None:
                joint_name = self.skeleton.bone_order_names[joint_idx]
                joint_visible = joint_name not in hidden_gizmo_joints
            cur_joint_gizmo = self.server.scene.add_transform_controls(
                f"/{self.name}/gizmo_joint_{joint_idx}",
                scale=0.075,
                line_width=4.0,
                active_axes=active_axes,
                disable_axes=disable_axes,
                disable_sliders=disable_sliders,
                disable_rotations=False,
                depth_test=False,  # render even when occluded
                position=self.joints_pos[self.cur_frame_idx, joint_idx].cpu().numpy(),
                wxyz=joints_wxyzs[joint_idx],
                visible=joint_visible,
            )
            self.joint_gizmos.append(cur_joint_gizmo)

            def set_callback_in_closure(i: int) -> None:
                @cur_joint_gizmo.on_update
                def _(_) -> None:
                    self.updating_joint_gizmos = True
                    new_local_joint_rots = self.joints_local_rot[self.cur_frame_idx].clone()
                    new_local_rot = tf.SO3(self.joint_gizmos[i].wxyz)
                    new_local_rot_mat_np = new_local_rot.as_matrix()
                    if joint_axis_indices is not None:
                        joint_name = self.skeleton.bone_order_names[i]
                        axis_idx = joint_axis_indices.get(joint_name)
                        if axis_idx is not None:
                            rotvec = new_local_rot.log()
                            axis = np.zeros(3, dtype=np.float64)
                            axis[axis_idx] = 1.0
                            new_local_rot_mat_np = tf.SO3.exp(rotvec[axis_idx] * axis).as_matrix()
                    new_local_rot_mat = torch.tensor(new_local_rot_mat_np).to(new_local_joint_rots.device)
                    new_local_joint_rots[i] = new_local_rot_mat

                    self.update_pose_at_frame(
                        self.cur_frame_idx,
                        joints_local_rot=new_local_joint_rots,
                    )

                    # handle root translation separately
                    cur_joints_pos = self.joints_pos[self.cur_frame_idx].clone()
                    if i == self.skeleton.root_idx:
                        new_root_pos = to_torch(
                            self.joint_gizmos[i].position,
                            device=self.joints_pos.device,
                        ).to(dtype=self.joints_pos.dtype)
                        root_diff = new_root_pos - self.joints_pos[self.cur_frame_idx, i]
                        if torch.norm(root_diff) > 1e-3:
                            # the root translation has been changed
                            # translate to gizmo position
                            cur_joints_pos += root_diff[None]
                            self.update_pose_at_frame(
                                self.cur_frame_idx,
                                joints_pos=cur_joints_pos,
                                joints_rot=self.joints_rot[self.cur_frame_idx],
                                joints_local_rot=self.joints_local_rot[self.cur_frame_idx],
                            )

                    self.updating_joint_gizmos = False

                    # update immediately to show user changes
                    self.set_frame(self.cur_frame_idx)

                    if i == self.skeleton.root_idx:
                        # update the 2D waypoint constraints as well if there is one
                        if "2D Root" in constraints:
                            root_2d_contraints = constraints["2D Root"]
                            # if there is a constraint at that frame, we want to update it
                            frame_idx = self.cur_frame_idx
                            if frame_idx in root_2d_contraints.keyframes:
                                new_root_pos[1] = 0.0  # force y to 0
                                for keyframe_id in root_2d_contraints.frame2keyid[frame_idx]:
                                    # add will modify the existing constraint
                                    root_2d_contraints.add_keyframe(
                                        keyframe_id,
                                        frame_idx,
                                        root_pos=new_root_pos,
                                        exists_ok=True,
                                    )

                    if "Full-Body" in constraints:
                        full_body_constraints = constraints["Full-Body"]
                        # if there is a constraint at that frame, we want to update it
                        frame_idx = self.cur_frame_idx
                        if frame_idx in full_body_constraints.keyframes:
                            for keyframe_id in full_body_constraints.frame2keyid[frame_idx]:
                                # add will modify the existing constraint
                                full_body_constraints.add_keyframe(
                                    keyframe_id,
                                    frame_idx,
                                    joints_pos=self.joints_pos[frame_idx],
                                    joints_rot=self.joints_rot[frame_idx],
                                    exists_ok=True,
                                )
                    if "End-Effectors" in constraints:
                        end_effector_constraints = constraints["End-Effectors"]
                        # if there is a constraint at that frame, we want to update it
                        frame_idx = self.cur_frame_idx
                        if frame_idx in end_effector_constraints.keyframes:
                            current_dict = end_effector_constraints.keyframes[frame_idx]
                            for keyframe_id, _ in end_effector_constraints.frame2keyid[frame_idx]:
                                # add will modify the existing constraint
                                end_effector_constraints.add_keyframe(
                                    keyframe_id,
                                    frame_idx,
                                    joints_pos=self.joints_pos[frame_idx],
                                    joints_rot=self.joints_rot[frame_idx],
                                    joint_names=current_dict["joint_names"],
                                    end_effector_type=current_dict["end_effector_type"],
                                    exists_ok=True,
                                )

            set_callback_in_closure(joint_idx)

    def clear_all_gizmos(self):
        self.updating_root_translation_gizmo = True
        self.updating_joint_gizmos = True
        if self.root_translation_gizmo is not None:
            self.server.scene.remove_by_name(self.root_translation_gizmo.name)
            self.root_translation_gizmo = None
        if self.joint_gizmos is not None:
            for joint_gizmo in self.joint_gizmos:
                self.server.scene.remove_by_name(joint_gizmo.name)
            self.joint_gizmos = None
        self.updating_root_translation_gizmo = False
        self.updating_joint_gizmos = False


#
# Constraint classes
#


class ConstraintSet:
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        display_name: Optional[str] = None,
    ):
        self.name = name
        self.server = server
        self.skeleton = skeleton
        self.display_name = display_name if display_name is not None else name

        self.keyframes = dict()  # frame_idx -> poses
        self.frame2keyid = dict()  # frame_idx -> list of keyframe ids at this frame
        self.scene_elements = dict()  # frame_idx -> meshes, labels, etc.
        self.interval_labels = dict()  # (start_frame_idx, end_frame_idx) -> interval_label
        self.labels_visible = True

    def set_label_visibility(self, visible: bool) -> None:
        """Show or hide constraint labels without deleting them."""
        self.labels_visible = visible
        for scene_data in self.scene_elements.values():
            label = scene_data.get("label")
            if label is not None:
                label.visible = visible
        for interval_label in self.interval_labels.values():
            interval_label.visible = visible

    def add_keyframe(self, keyframe_id: str, frame_idx: int, pose_data: torch.Tensor):
        """Adds a single keyframe at the given frame with the given pose data.

        Args:
            keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx.
            frame_idx: int, frame index to add the keyframe at
            pose_data: torch.Tensor, e.g. full-body pose, EE pose, 2D root pose, etc.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def add_interval(
        self,
        interval_id: str,
        start_frame_idx: int,
        end_frame_idx: int,
        pose_seq_data: torch.Tensor,
    ):
        """Adds a keyframe interval between the given start and end frames with the given pose data.

        Args:
            interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx.
            start_frame_idx: int, start frame index of the interval
            end_frame_idx: int, end frame index of the interval
            pose_seq_data: torch.Tensor, data for constrained interval, e.g. full-body poses, EE poses, 2D root poses, etc.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def _add_interval_label(self, start_frame_idx: int, end_frame_idx: int):
        """
        Adds an interval label between the given start and end frames
        Args:
            start_frame_idx: int, start frame index of the interval
            end_frame_idx: int, end frame index of the interval
        """
        mid = int((start_frame_idx + end_frame_idx) / 2)
        interval_label_pos = self._get_label_pos(mid)
        interval_label = self.server.scene.add_label(
            name=f"/{self.name}/interval_label_{start_frame_idx}_{end_frame_idx}",
            text=f"{self.display_name} @ [{start_frame_idx}, {end_frame_idx}]",
            position=interval_label_pos,
            font_size_mode="screen",
            font_screen_scale=0.7,
            anchor="center-center",
        )
        interval_label.visible = self.labels_visible
        self.interval_labels[(start_frame_idx, end_frame_idx)] = interval_label

    def remove_keyframe(self, keyframe_id: str, frame_idx: int):
        """
        Removes a keyframe at the given frame
        Args:
            keyframe_id: str, id for the keyframe to remove
            frame_idx: int, frame index to remove the keyframe at
        """
        raise NotImplementedError("Subclasses must implement this method")

    def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int):
        """
        Removes an interval between the given start and end frames
        Args:
            interval_id: str, id for the interval to remove
            start_frame_idx: int, start frame index of the interval
            end_frame_idx: int, end frame index of the interval
        """
        raise NotImplementedError("Subclasses must implement this method")

    def _get_label_pos(self, frame_idx: int):
        """
        Returns the position of where to place the displayed label for the given frame index
        Args:
            frame_idx: int, frame index to get the label position for
        """
        raise NotImplementedError("Subclasses must implement this method")

    def _remove_interval_and_update_label(self, interval_id: str, start_frame_idx: int, end_frame_idx: int):
        """
        Removes an interval between the given start and end frames and updates the label
        Args:
            start_frame_idx: int, start frame index of the interval
            end_frame_idx: int, end frame index of the interval
        """
        for frame_idx in range(start_frame_idx, end_frame_idx + 1):
            self.remove_keyframe(interval_id, frame_idx)

        # Update interval labels that overlap with the removed range
        intervals_to_update = []
        for (interval_start, interval_end), label in list(self.interval_labels.items()):
            # Check if intervals overlap
            if interval_start <= end_frame_idx and interval_end >= start_frame_idx:
                intervals_to_update.append((interval_start, interval_end, label))

        for interval_start, interval_end, label in intervals_to_update:
            # Remove old label from scene and dict
            self.server.scene.remove_by_name(label.name)
            del self.interval_labels[(interval_start, interval_end)]

            new_start, new_end = update_interval(interval_start, interval_end, start_frame_idx, end_frame_idx)

            if new_start is None or new_end is None:
                continue

            # Create updated label with new range
            if new_start <= new_end:
                # Position label at midpoint - these keyframes are guaranteed to exist
                # since the new range is outside the removal range
                mid_frame = (new_start + new_end) // 2
                label_pos = self._get_label_pos(mid_frame)
                new_label = self.server.scene.add_label(
                    name=f"/{self.name}/interval_label_{new_start}_{new_end}",
                    text=f"{self.display_name} @ [{new_start}, {new_end}]",
                    position=label_pos,
                    font_size_mode="screen",
                    font_screen_scale=0.7,
                    anchor="center-center",
                )
                new_label.visible = self.labels_visible
                self.interval_labels[(new_start, new_end)] = new_label

    def get_constraint_info(self, device: Optional[str] = None):
        """Returns constraint information for generation (torch) or UI (numpy)."""
        raise NotImplementedError("Subclasses must implement this method")

    def get_frame_idx(self):
        """Returns all constrained frame indices in the set."""
        return [frame_idx for frame_idx in list(self.keyframes.keys())]

    def set_keyframe_visibility(self, keyframe_idx: int, visible: bool):
        """Sets the visibility of scene elements at the given keyframe index.

        Args:
            keyframe_idx: int, keyframe index to set visibility for
            visible: bool, whether to make the elements visible
        """
        raise NotImplementedError("Subclasses must implement this method")

    def clear(self, frame_idx: Optional[int] = None):
        """
        Clears all keyframes and intervals from the constraint set
        Args:
            frame_idx: int, sing frame index to clear if given
        """
        raise NotImplementedError("Subclasses must implement this method")


def build_constraint_set_table_markdown(constraint_list: List[ConstraintSet]):
    markdown = "| Track | Frame Num |\n"
    markdown += "|------|----------|\n"

    # Sort constraints by frame_idx
    for constraint in constraint_list:
        frame_info = constraint.get_frame_idx()
        if len(frame_info) > 0:
            frame_info = ", ".join([str(frame) for frame in sorted(frame_info)])
        else:
            frame_info = "-"
        markdown += f"| {constraint.display_name} | {frame_info} |\n"

    return markdown


def update_interval(interval_start, interval_end, start_frame_idx, end_frame_idx):
    """Updates an interval after removing the range from start_frame_idx to end_frame_idx."""
    # Calculate new range after removing [start_frame_idx, end_frame_idx]
    # Case 1: Removal fully contains the interval -> delete entirely
    if start_frame_idx <= interval_start and end_frame_idx >= interval_end:
        return None, None  # Already removed, don't recreate
    # Case 2: Removal is at the start of interval -> shrink from start
    elif start_frame_idx <= interval_start and end_frame_idx < interval_end:
        new_start = end_frame_idx + 1
        new_end = interval_end
    # Case 3: Removal is at the end of interval -> shrink from end
    elif start_frame_idx > interval_start and end_frame_idx >= interval_end:
        new_start = interval_start
        new_end = start_frame_idx - 1
    # Case 4: Removal is in the middle -> keep the larger portion
    else:  # start_frame_idx > interval_start and end_frame_idx < interval_end
        left_size = start_frame_idx - interval_start
        right_size = interval_end - end_frame_idx
        if left_size >= right_size:
            new_start = interval_start
            new_end = start_frame_idx - 1
        else:
            new_start = end_frame_idx + 1
            new_end = interval_end
    return new_start, new_end


class FullbodyKeyframeSet(ConstraintSet):
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        display_name: Optional[str] = None,
    ):
        super().__init__(name, server, skeleton, display_name=display_name)

    def add_keyframe(
        self,
        keyframe_id: str,
        frame_idx: int,
        joints_pos: torch.Tensor | np.ndarray,
        joints_rot: torch.Tensor | np.ndarray,
        viz_label: bool = True,
        exists_ok: bool = False,
    ):
        """Adds a single full-body keyframe at the given frame or updates the existing one at this
        frame. Note if a keyframe already exists at this frame, it will be updated to the given
        pose.

        Args:
            keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx.
            frame_idx: int, frame index to add the keyframe at
            joints_pos: torch.Tensor, [J, 3] joints positions to add the keyframe at
        """
        # create/update scene elements
        if frame_idx in self.keyframes:
            skeleton_mesh = self.scene_elements[frame_idx]["skeleton_mesh"]
            skeleton_mesh.set_pose(to_torch(joints_pos))
            if viz_label and "label" in self.scene_elements[frame_idx]:
                label = self.scene_elements[frame_idx]["label"]
                label.position = to_numpy(joints_pos)[self.skeleton.root_idx]
                label.visible = self.labels_visible
        else:
            # create skeleton to visualize the full-body constraint
            skeleton_mesh = SkeletonMesh(
                f"/{self.name}/skeleton_{frame_idx}",
                self.server,
                self.skeleton,
                joint_color=(255, 235, 0),
                bone_color=(255, 0, 0),
                starting_joints_pos=to_torch(joints_pos),
            )
            self.scene_elements[frame_idx] = {
                "skeleton_mesh": skeleton_mesh,
            }
            if viz_label:
                label = self.server.scene.add_label(
                    name=f"/{self.name}/label_{frame_idx}",
                    text=f"{self.display_name} @ {frame_idx}",
                    position=to_numpy(joints_pos)[self.skeleton.root_idx],
                    font_size_mode="screen",
                    font_screen_scale=0.7,
                    anchor="center-center",
                )
                label.visible = self.labels_visible
                self.scene_elements[frame_idx]["label"] = label

        # set/update data
        self.keyframes[frame_idx] = {
            "joints_pos": to_numpy(joints_pos),
            "joints_rot": to_numpy(joints_rot),
        }

        if frame_idx not in self.frame2keyid:
            self.frame2keyid[frame_idx] = []

        if keyframe_id in self.frame2keyid[frame_idx]:
            if not exists_ok:
                raise AssertionError("keyframe_id already exists in this frame!")
        else:
            self.frame2keyid[frame_idx].append(keyframe_id)

    def add_interval(
        self,
        interval_id: str,
        start_frame_idx: int,
        end_frame_idx: int,
        joints_pos: torch.Tensor,
        joints_rot: torch.Tensor,
    ):
        """Adds a full-body keyframe interval between the given start and end frames.

        Args:
            start_frame_idx: int, start frame index of the interval
            end_frame_idx: int, end frame index of the interval
            joints_pos: torch.Tensor, [T, J, 3] joints positions within the interval
        """
        assert joints_pos.shape[0] == end_frame_idx - start_frame_idx + 1
        for frame_idx in range(start_frame_idx, end_frame_idx + 1):
            rel_idx = frame_idx - start_frame_idx
            self.add_keyframe(
                interval_id,
                frame_idx,
                joints_pos[rel_idx],
                joints_rot[rel_idx],
                viz_label=False,
            )

        # add separate interval label
        self._add_interval_label(start_frame_idx, end_frame_idx)

    def remove_keyframe(self, keyframe_id: str, frame_idx: int):
        if frame_idx not in self.keyframes:
            return
        if keyframe_id not in self.frame2keyid[frame_idx]:
            return
        self.frame2keyid[frame_idx].remove(keyframe_id)
        if len(self.frame2keyid[frame_idx]) == 0:
            del self.frame2keyid[frame_idx]
            self.clear(frame_idx)

    def _get_label_pos(self, frame_idx: int):
        return self.keyframes[frame_idx]["joints_pos"][self.skeleton.root_idx]

    def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int):
        self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx)

    def get_constraint_info(self, device: Optional[str] = None):
        all_joints_pos = []
        all_joints_rot = []
        for v in self.keyframes.values():
            joints_pos = to_torch(v["joints_pos"], device=device)
            joints_rot = to_torch(v["joints_rot"], device=device)
            if len(joints_pos.shape) == 2:
                all_joints_pos.append(joints_pos[None])
            else:
                all_joints_pos.append(joints_pos)
            if len(joints_rot.shape) == 3:
                all_joints_rot.append(joints_rot[None])
            else:
                all_joints_rot.append(joints_rot)

        all_joints_pos = torch.cat(all_joints_pos, dim=0) if len(all_joints_pos) > 0 else None
        all_joints_rot = torch.cat(all_joints_rot, dim=0) if len(all_joints_rot) > 0 else None

        return {
            "frame_idx": self.get_frame_idx(),
            "joints_pos": all_joints_pos,
            "joints_rot": all_joints_rot,
        }

    def set_keyframe_visibility(self, keyframe_idx: int, visible: bool):
        """Sets the visibility of scene elements at the given keyframe index."""
        if keyframe_idx not in self.scene_elements:
            return

        scene_elements = self.scene_elements[keyframe_idx]
        if "skeleton_mesh" in scene_elements:
            skeleton_mesh = scene_elements["skeleton_mesh"]
            if hasattr(skeleton_mesh, "joints_batched_mesh"):
                skeleton_mesh.joints_batched_mesh.visible = visible
            if hasattr(skeleton_mesh, "bones_batched_mesh"):
                skeleton_mesh.bones_batched_mesh.visible = visible

        if "label" in scene_elements:
            label = scene_elements["label"]
            if hasattr(label, "visible"):
                label.visible = visible

    def clear(self, frame_idx: Optional[int] = None):
        frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx]
        for fidx in frame_idx_list:
            self.scene_elements[fidx]["skeleton_mesh"].clear()
            if "ee_rotation_axes" in self.scene_elements[fidx]:
                self.server.scene.remove_by_name(self.scene_elements[fidx]["ee_rotation_axes"].name)
            if "label" in self.scene_elements[fidx]:
                self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name)

            self.keyframes.pop(fidx)
            self.scene_elements.pop(fidx)

        if frame_idx is None:
            # clear all interval labels if clearing all keyframes
            for interval_label in list(self.interval_labels.values()):
                self.server.scene.remove_by_name(interval_label.name)
            self.interval_labels.clear()


class EEJointsKeyframeSet(ConstraintSet):
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        display_name: Optional[str] = None,
    ):
        super().__init__(name, server, skeleton, display_name=display_name)

        # frame_idx -> list of (keyframe_id, joint_names) at this frame
        self.frame2keyid = dict()

    def create_scene_elements(
        self,
        frame_idx: int,
        joints_pos: torch.Tensor | np.ndarray,
        joints_rot: Optional[torch.Tensor | np.ndarray],
        joint_names: List[str],
        viz_label: bool = True,
    ):
        # create skeleton to visualize the full-body constraint
        ee_joint_indices = []
        ee_gizmo_indices = []
        constrained_bone_idx = []
        for joint_name in joint_names:
            if joint_name == "Hips":
                continue
            elif joint_name in ["LeftHand", "RightHand", "LeftFoot", "RightFoot"]:
                expanded_joint_names = {
                    "LeftHand": self.skeleton.left_hand_joint_names,
                    "RightHand": self.skeleton.right_hand_joint_names,
                    "LeftFoot": self.skeleton.left_foot_joint_names,
                    "RightFoot": self.skeleton.right_foot_joint_names,
                }[joint_name]
                ee_joint_indices.extend([self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names])
                if len(expanded_joint_names) > 1:
                    ee_gizmo_indices.extend(
                        [self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names[:-1]]
                    )
                constrained_bone_idx.extend(
                    [self.skeleton.bone_order_names_index[joint] - 1 for joint in expanded_joint_names[1:]]
                )
            else:
                raise ValueError(f"Invalid joint name: {joint_name}")

        # de-duplicate while preserving order
        ee_joint_indices = list(dict.fromkeys(ee_joint_indices))
        ee_gizmo_indices = list(dict.fromkeys(ee_gizmo_indices))
        constrained_bone_idx = list(dict.fromkeys(constrained_bone_idx))

        constrained_idx = np.array([self.skeleton.root_idx] + ee_joint_indices, dtype=np.intp)
        constrained_bone_idx = np.array(constrained_bone_idx, dtype=np.intp)

        # create skeleton to visualize the full-body constraint
        joint_color = np.full((self.skeleton.nbjoints, 3), (220, 220, 220))
        bone_color = np.full((self.skeleton.nbjoints - 1, 3), (220, 220, 220))
        # color constrained joints and bones red
        joint_color[constrained_idx] = (255, 0, 0)
        if len(constrained_bone_idx) > 0:
            bone_color[constrained_bone_idx] = (255, 0, 0)
        skeleton_mesh = SkeletonMesh(
            f"/{self.name}/skeleton_{frame_idx}",
            self.server,
            self.skeleton,
            joint_color=joint_color,
            bone_color=bone_color,
            starting_joints_pos=to_torch(joints_pos),
        )

        self.scene_elements[frame_idx] = {
            "skeleton_mesh": skeleton_mesh,
        }
        joints_pos_np = to_numpy(joints_pos)
        joints_rot_np = to_numpy(joints_rot) if joints_rot is not None else None
        if joints_rot_np is not None and len(ee_gizmo_indices) > 0:
            ee_axes = self.server.scene.add_batched_axes(
                f"/{self.name}/ee_rot_axes_{frame_idx}",
                batched_wxyzs=tf.SO3.from_matrix(joints_rot_np[ee_gizmo_indices]).wxyz,
                batched_positions=joints_pos_np[ee_gizmo_indices],
                axes_length=0.07,
                axes_radius=0.007,
            )
            self.scene_elements[frame_idx]["ee_rotation_axes"] = ee_axes
        if viz_label:
            label = self.server.scene.add_label(
                name=f"/{self.name}/label_{frame_idx}",
                text=f"{self.display_name} @ {frame_idx}",
                position=joints_pos_np[self.skeleton.root_idx] + np.array([0.0, 0.05, 0.0]),
                font_size_mode="screen",
                font_screen_scale=0.7,
                anchor="bottom-center",
            )
            label.visible = self.labels_visible
            self.scene_elements[frame_idx]["label"] = label

    def add_keyframe(
        self,
        keyframe_id: str,
        frame_idx: int,
        joints_pos: torch.Tensor | np.ndarray,
        joints_rot: torch.Tensor | np.ndarray,
        joint_names: List[str],
        end_effector_type: str,
        viz_label: bool = True,
        exists_ok: bool = False,
    ):
        """Adds a single EE keyframe at the given frame or updates the existing one at this frame.

        Args:
            keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx.
            frame_idx: int, frame index to add the keyframe at
            joints_pos: torch.Tensor, [J, 3] joints positions to add the keyframe at
            joints_rot: torch.Tensor, [J, 3, 3] joints rotation matrices to add the keyframe at
            joint_names: List[str], names of the joints to add the keyframe at
        """
        need_create_viz = True
        joint_names_input = joint_names

        if not isinstance(end_effector_type, set):
            end_effector_type = set([end_effector_type])

        # create/update scene elements
        if frame_idx in self.keyframes:
            if joint_names != self.keyframes[frame_idx]["joint_names"]:
                # merge together with existing constraint if needed
                joint_names = set(joint_names)
                joint_names.update(set(self.keyframes[frame_idx]["joint_names"]))
                joint_names = list(joint_names)
                end_effector_type.update(self.keyframes[frame_idx]["end_effector_type"])
                # need to re-create viz elements
                self.clear(frame_idx)
            else:
                need_create_viz = False
                # overwrite the pose with the latest one
                skeleton_mesh = self.scene_elements[frame_idx]["skeleton_mesh"]
                skeleton_mesh.set_pose(to_torch(joints_pos))
                if "ee_rotation_axes" in self.scene_elements[frame_idx]:
                    ee_gizmo_indices = []
                    for joint_name in joint_names:
                        if joint_name == "Hips":
                            continue
                        elif joint_name in [
                            "LeftHand",
                            "RightHand",
                            "LeftFoot",
                            "RightFoot",
                        ]:
                            expanded_joint_names = {
                                "LeftHand": self.skeleton.left_hand_joint_names,
                                "RightHand": self.skeleton.right_hand_joint_names,
                                "LeftFoot": self.skeleton.left_foot_joint_names,
                                "RightFoot": self.skeleton.right_foot_joint_names,
                            }[joint_name]
                            if len(expanded_joint_names) > 1:
                                ee_gizmo_indices.extend(
                                    [self.skeleton.bone_order_names_index[joint] for joint in expanded_joint_names[:-1]]
                                )
                        else:
                            raise ValueError(f"Invalid joint name: {joint_name}")
                    ee_gizmo_indices = list(dict.fromkeys(ee_gizmo_indices))
                    if len(ee_gizmo_indices) > 0:
                        ee_axes = self.scene_elements[frame_idx]["ee_rotation_axes"]
                        joints_pos_np = to_numpy(joints_pos)
                        joints_rot_np = to_numpy(joints_rot)
                        ee_axes.batched_positions = joints_pos_np[ee_gizmo_indices]
                        ee_axes.batched_wxyzs = tf.SO3.from_matrix(joints_rot_np[ee_gizmo_indices]).wxyz
                if viz_label and "label" in self.scene_elements[frame_idx]:
                    label = self.scene_elements[frame_idx]["label"]
                    label.position = to_numpy(joints_pos)[self.skeleton.root_idx]
                    label.visible = self.labels_visible

        if need_create_viz:
            self.create_scene_elements(frame_idx, joints_pos, joints_rot, joint_names, viz_label=viz_label)

        # set/update data
        self.keyframes[frame_idx] = {
            "joints_pos": to_numpy(joints_pos),
            "joints_rot": to_numpy(joints_rot),
            "joint_names": joint_names,
            "end_effector_type": end_effector_type,
        }

        if frame_idx not in self.frame2keyid:
            self.frame2keyid[frame_idx] = []

        known_keyframe_ids = {k: idx for idx, (k, _) in enumerate(self.frame2keyid[frame_idx])}

        if keyframe_id in known_keyframe_ids.keys():
            if not exists_ok:
                raise AssertionError("keyframe_id already exists in this frame!")
            idx = known_keyframe_ids[keyframe_id]
            # override previous exisiting keyframe
            self.frame2keyid[frame_idx][idx] = (keyframe_id, joint_names_input)
        else:
            # track which subset of joints are constrained by this keyframe_id
            self.frame2keyid[frame_idx].append((keyframe_id, joint_names_input))

    def add_interval(
        self,
        interval_id: str,
        start_frame_idx: int,
        end_frame_idx: int,
        joints_pos: torch.Tensor | np.ndarray,
        joints_rot: torch.Tensor | np.ndarray,
        joint_names: List[str],
        end_effector_type: str,
    ):
        """Adds an interval of EE keyframes at the given frame or updates the existing one at this
        frame.

        Args:
            interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx.
            start_frame_idx: int, start frame index to add the interval at
            end_frame_idx: int, end frame index to add the interval at
            joints_pos: torch.Tensor, [T, J, 3] joints positions to add the interval at
            joints_rot: torch.Tensor, [T, J, 3, 3] joints rotation matrices to add the interval at
            joint_names: List[str], names of the joints to add for the entire interval
        """
        num_frames = end_frame_idx - start_frame_idx + 1
        joints_pos_np = to_numpy(joints_pos)
        joints_rot_np = to_numpy(joints_rot)
        assert joints_pos_np.shape[0] == num_frames
        assert joints_rot_np.shape[0] == num_frames

        for frame_idx in range(start_frame_idx, end_frame_idx + 1):
            rel_idx = frame_idx - start_frame_idx
            self.add_keyframe(
                interval_id,
                frame_idx,
                joints_pos_np[rel_idx],
                joints_rot_np[rel_idx],
                joint_names,
                end_effector_type,
                viz_label=False,
            )
        self._add_interval_label(start_frame_idx, end_frame_idx)

    def remove_keyframe(self, keyframe_id: str, frame_idx: int):
        """Removes a keyframe at the given frame or updates the existing one at this frame by
        removing the specified joints.

        Args:
            keyframe_id: str, id for the keyframe to remove. This determines which joints to remove.
            frame_idx: int, frame index to remove the keyframe at
        """
        if frame_idx not in self.keyframes:
            return

        remaining_joint_names = set()
        delete_idx = None
        for i, (keyid, joint_names) in enumerate(self.frame2keyid[frame_idx]):
            if keyid == keyframe_id:
                delete_idx = i
            else:
                remaining_joint_names.update(joint_names)
        if delete_idx is None:
            # this keyframe_id is not in the specified frame
            return

        self.frame2keyid[frame_idx].pop(delete_idx)
        if len(remaining_joint_names) == 0:
            # no more keyframes in this frame, clear the frame
            del self.frame2keyid[frame_idx]
            self.clear(frame_idx)
            return

        # only deleting part of keyframe (potentially some subset of joints)
        # delete the old visualization and add a new one with the updated joint set
        new_joint_names = list(remaining_joint_names)
        self.clear(frame_idx, scene_elements_only=True)
        joints_pos = self.keyframes[frame_idx]["joints_pos"]
        joints_rot = self.keyframes[frame_idx]["joints_rot"]
        self.create_scene_elements(frame_idx, joints_pos, joints_rot, new_joint_names)
        self.keyframes[frame_idx]["joint_names"] = new_joint_names

    def _get_label_pos(self, frame_idx: int):
        return self.keyframes[frame_idx]["joints_pos"][self.skeleton.root_idx]

    def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int):
        self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx)

    def get_constraint_info(self, device: Optional[str] = None):
        all_joints_pos = []
        all_joints_rot = []
        all_joints_names = []
        all_end_effector_type = []
        for v in self.keyframes.values():
            joints_pos = to_torch(v["joints_pos"], device=device)
            joints_rot = to_torch(v["joints_rot"], device=device)
            if len(joints_pos.shape) == 2:
                all_joints_pos.append(joints_pos[None])
            else:
                all_joints_pos.append(joints_pos)
            if len(joints_rot.shape) == 3:
                all_joints_rot.append(joints_rot[None])
            else:
                all_joints_rot.append(joints_rot)
            all_joints_names.append(v["joint_names"])
            all_end_effector_type.append(v["end_effector_type"])

        all_joints_pos = torch.cat(all_joints_pos, dim=0) if len(all_joints_pos) > 0 else None
        all_joints_rot = torch.cat(all_joints_rot, dim=0) if len(all_joints_rot) > 0 else None

        return {
            "frame_idx": self.get_frame_idx(),
            "joints_pos": all_joints_pos,
            "joints_rot": all_joints_rot,
            "joint_names": all_joints_names,
            "end_effector_type": all_end_effector_type,
        }

    def set_keyframe_visibility(self, keyframe_idx: int, visible: bool, show_rotation_axes: bool = True):
        """Sets the visibility of scene elements at the given keyframe index."""
        if keyframe_idx not in self.scene_elements:
            return

        scene_elements = self.scene_elements[keyframe_idx]
        if "skeleton_mesh" in scene_elements:
            skeleton_mesh = scene_elements["skeleton_mesh"]
            if hasattr(skeleton_mesh, "joints_batched_mesh"):
                skeleton_mesh.joints_batched_mesh.visible = visible
            if hasattr(skeleton_mesh, "bones_batched_mesh"):
                skeleton_mesh.bones_batched_mesh.visible = visible

        if "ee_rotation_axes" in scene_elements:
            ee_axes = scene_elements["ee_rotation_axes"]
            if hasattr(ee_axes, "visible"):
                ee_axes.visible = visible and show_rotation_axes

        if "label" in scene_elements:
            label = scene_elements["label"]
            if hasattr(label, "visible"):
                label.visible = visible

    def clear(self, frame_idx: Optional[int] = None, scene_elements_only: bool = False):
        frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx]
        for fidx in frame_idx_list:
            self.scene_elements[fidx]["skeleton_mesh"].clear()
            if "ee_rotation_axes" in self.scene_elements[fidx]:
                self.server.scene.remove_by_name(self.scene_elements[fidx]["ee_rotation_axes"].name)
            if "label" in self.scene_elements[fidx]:
                self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name)
            self.scene_elements.pop(fidx)
            if not scene_elements_only:
                self.keyframes.pop(fidx)

        if frame_idx is None:
            # clear all interval labels if clearing all keyframes
            for interval_label in list(self.interval_labels.values()):
                self.server.scene.remove_by_name(interval_label.name)
            self.interval_labels.clear()


class RootKeyframe2DSet(ConstraintSet):
    def __init__(
        self,
        name: str,
        server: viser.ViserServer,
        skeleton: SkeletonBase,
        display_name: Optional[str] = None,
    ):
        super().__init__(name, server, skeleton, display_name=display_name)
        self.dense_path = False
        self.smooth_path = True
        self.line_segments = None  # visualization of dense path
        # Cache for interpolated path
        self._cached_t = None
        self._cached_path3d = None
        # Store root headings for each keyframe
        self.root_headings = {}  # frame_idx -> heading (float)

    def add_keyframe(
        self,
        keyframe_id: str,
        frame_idx: int,
        root_pos: torch.Tensor | np.ndarray,
        global_root_heading: Optional[float] = None,
        viz_label: bool = True,
        update_path: bool = True,
        exists_ok: bool = False,
        add_annulus: bool = True,
    ):
        """Adds a single 2D root keyframe at the given frame or updates the existing one at this
        frame.

        Args:
            keyframe_id: str, id for the keyframe. Must be unique within the given frame_idx.
            frame_idx: int, frame index to add the keyframe at
            root_pos: torch.Tensor, [3] root position to add the keyframe at, y entry (index 1) should be 0
            viz_label: bool, whether to visualize the label for the keyframe
        """
        root_pos_np = to_numpy(root_pos)

        # Convert heading angle to 2D direction vector for visualization
        heading_2d = None
        if global_root_heading is not None:
            # Heading is in radians, convert to [x, z] direction vector as [cos, sin]
            heading_2d = np.array([np.cos(global_root_heading), np.sin(global_root_heading)])

        if frame_idx in self.keyframes:
            waypoint = self.scene_elements[frame_idx]["waypoint"]
            waypoint.update_position(root_pos_np, heading=heading_2d)
            if viz_label and "label" in self.scene_elements[frame_idx]:
                label = self.scene_elements[frame_idx]["label"]
                label.position = root_pos.cpu().numpy()
                label.visible = self.labels_visible
        else:
            waypoint = WaypointMesh(
                f"/{self.name}/{keyframe_id}",
                self.server,
                position=root_pos_np,
                heading=heading_2d,
                add_annulus=add_annulus,
            )
            self.scene_elements[frame_idx] = {
                "waypoint": waypoint,
            }
            if viz_label:
                label = self.server.scene.add_label(
                    name=f"/{self.name}/label_{frame_idx}",
                    text=f"{self.display_name} @ {frame_idx}",
                    position=root_pos_np,
                    font_size_mode="screen",
                    font_screen_scale=0.7,
                    anchor="bottom-left",
                )
                label.visible = self.labels_visible
                self.scene_elements[frame_idx]["label"] = label

        # set/update data
        self.keyframes[frame_idx] = root_pos_np
        if global_root_heading is not None:
            self.root_headings[frame_idx] = global_root_heading
        if frame_idx not in self.frame2keyid:
            self.frame2keyid[frame_idx] = []

        if keyframe_id in self.frame2keyid[frame_idx]:
            if not exists_ok:
                raise AssertionError("keyframe_id already exists in this frame!")
        else:
            self.frame2keyid[frame_idx].append(keyframe_id)

        # need to update path visualization
        if self.dense_path and self.line_segments is None:
            # visualize dense path with line segments
            self.line_segments = self.server.scene.add_line_segments(
                name=f"/{self.name}/line_segments",
                points=np.zeros((1, 2, 3)),
                colors=(255, 0, 0),
                line_width=5.0,
            )
        if update_path:
            self.update_line_segments()

    def add_interval(
        self,
        interval_id: str,
        start_frame_idx: int,
        end_frame_idx: int,
        root_pos: torch.Tensor | np.ndarray,
        add_annulus: bool = True,
    ):
        """Adds an interval of 2D root keyframes between the given start and end frames.

        Args:
            interval_id: str, id for the interval. Must be unique within the given start_frame_idx and end_frame_idx.
            start_frame_idx: int, start frame index to add the interval at
            end_frame_idx: int, end frame index to add the interval at
            root_pos: torch.Tensor, [T, 3] root positions to add the interval at
        """
        root_pos_np = to_numpy(root_pos)
        assert root_pos_np.shape[0] == end_frame_idx - start_frame_idx + 1
        for frame_idx in range(start_frame_idx, end_frame_idx + 1):
            rel_idx = frame_idx - start_frame_idx
            self.add_keyframe(
                f"{interval_id}/{frame_idx}",
                frame_idx,
                root_pos_np[rel_idx],
                viz_label=False,
                update_path=False,
                add_annulus=add_annulus,
            )
        self._add_interval_label(start_frame_idx, end_frame_idx)
        if self.line_segments is not None:
            self.update_line_segments()

    def set_smooth_path(self, smooth_path: bool):
        self.smooth_path = smooth_path
        if self.line_segments is not None:
            self.update_line_segments()

    def set_dense_path(self, dense_path: bool):
        """If dense_path is True, will make the path dense by interpolated between added keyframes.

        Args:
            dense_path: bool, whether to make the path dense
        """
        self.dense_path = dense_path
        if self.dense_path:
            # visualize dense path with line segments
            self.line_segments = self.server.scene.add_line_segments(
                name=f"/{self.name}/line_segments",
                points=np.zeros((1, 2, 3)),
                colors=(255, 0, 0),
                line_width=5.0,
            )
            self.update_line_segments()
        else:
            if self.line_segments is not None:
                self.server.scene.remove_by_name(self.line_segments.name)
                self.line_segments = None

    # TODO: accelerate interp1d. Note that np and torch interp gives different results, much less smooth
    def interpolate_path(self, t: np.ndarray):
        """Interpolates the path between the given frame indices.

        Args:
            t: np.ndarray, frame indices to interpolate at
        """
        cur_info = self._get_sparse_constraint_info()
        frame_idx = cur_info["frame_idx"]
        all_root_pos = cur_info["root_pos"]

        x = all_root_pos[:, 0]
        z = all_root_pos[:, 2]

        kind = "linear"
        # if self.smooth_path and len(frame_idx) >= 3:
        # kind = "quadratic"

        interp_x = interp1d(frame_idx, x, kind=kind)
        interp_z = interp1d(frame_idx, z, kind=kind)

        x_new = interp_x(t)
        z_new = interp_z(t)

        path3d = np.stack([x_new, np.zeros_like(x_new), z_new], axis=1)

        if self.smooth_path and len(frame_idx) >= 3:
            start_time = time.time()
            path3d = get_smooth_root_pos(torch.from_numpy(path3d[None]))[0].numpy()
            print(f"Time taken to smooth path: {time.time() - start_time} seconds")
        return path3d

    def update_line_segments(self, frame_idx: int = 0):
        if len(self.keyframes) < 2:
            return

        t = np.array(sorted(self.get_frame_idx()))
        if self.smooth_path:
            # more points for smoothed curve
            # t = np.linspace(t[0], t[-1], 100)
            t = np.arange(max(t[0], frame_idx), t[-1] + 1)

        path3d = self.interpolate_path(t)

        # Cache the computed t and path3d for use in get_constraint_info
        self._cached_t = t
        self._cached_path3d = path3d

        points = np.zeros((len(path3d) - 1, 2, 3))
        points[:, 0] = path3d[:-1]
        points[:, 1] = path3d[1:]

        self.line_segments.points = points

    def remove_keyframe(self, keyframe_id: str, frame_idx: int):
        if frame_idx not in self.keyframes:
            return
        if keyframe_id not in self.frame2keyid[frame_idx]:
            return
        self.frame2keyid[frame_idx].remove(keyframe_id)
        if len(self.frame2keyid[frame_idx]) == 0:
            del self.frame2keyid[frame_idx]
            # Also remove heading if it exists
            if frame_idx in self.root_headings:
                del self.root_headings[frame_idx]
            self.clear(frame_idx)
            if self.line_segments is not None:
                self.update_line_segments()

    def _get_label_pos(self, frame_idx: int):
        return self.keyframes[frame_idx]

    def remove_interval(self, interval_id: str, start_frame_idx: int, end_frame_idx: int):
        self._remove_interval_and_update_label(interval_id, start_frame_idx, end_frame_idx)

    def _get_sparse_constraint_info(self):
        all_root_pos = []
        all_root_headings = []
        frame_indices = self.get_frame_idx()

        for frame_idx in frame_indices:
            v = self.keyframes[frame_idx]
            # Handle both numpy arrays and torch tensors
            if isinstance(v, torch.Tensor):
                if len(v.shape) == 1:
                    all_root_pos.append(v.unsqueeze(0).cpu().numpy())
                else:
                    all_root_pos.append(v.cpu().numpy())
            else:  # numpy array
                if len(v.shape) == 1:
                    all_root_pos.append(v[np.newaxis, :])
                else:
                    all_root_pos.append(v)
            # Get heading if it exists
            if frame_idx in self.root_headings:
                all_root_headings.append(self.root_headings[frame_idx])

        if len(all_root_pos) > 0:
            all_root_pos = np.concatenate(all_root_pos, axis=0)
        else:
            all_root_pos = None

        result = {
            "frame_idx": frame_indices,
            "root_pos": all_root_pos,
        }

        if len(all_root_headings) > 0 and len(all_root_headings) == len(frame_indices):
            result["global_root_heading"] = torch.tensor(all_root_headings, dtype=torch.float32)

        return result

    def get_constraint_info(self):
        if not self.dense_path or len(self.keyframes) == 0:
            return self._get_sparse_constraint_info()
        else:
            # Use cached path if available, otherwise compute it
            if self._cached_t is not None and self._cached_path3d is not None:
                t = self._cached_t
                path3d = self._cached_path3d
            else:
                frame_idx_list = self.get_frame_idx()
                min_frame_idx = min(frame_idx_list)
                max_frame_idx = max(frame_idx_list)
                t = np.arange(min_frame_idx, max_frame_idx + 1)
                path3d = self.interpolate_path(t)

            result = {
                "frame_idx": t.tolist(),
                "root_pos": torch.tensor(path3d, dtype=torch.float32),
            }

            # Interpolate headings if any exist
            if len(self.root_headings) > 0:
                frame_idx_list = self.get_frame_idx()
                # Check if we have headings for all keyframes
                if all(fi in self.root_headings for fi in frame_idx_list):
                    # Interpolate headings for dense path
                    heading_values = [self.root_headings[fi] for fi in sorted(frame_idx_list)]
                    heading_frames = sorted(frame_idx_list)

                    # Linear interpolation of headings
                    interpolated_headings = []
                    for frame in t:
                        if frame in self.root_headings:
                            interpolated_headings.append(self.root_headings[frame])
                        else:
                            # Find surrounding keyframes and interpolate
                            left_frame = max([f for f in heading_frames if f <= frame], default=None)
                            right_frame = min([f for f in heading_frames if f >= frame], default=None)

                            if left_frame is not None and right_frame is not None and left_frame != right_frame:
                                # Linear interpolation
                                alpha = (frame - left_frame) / (right_frame - left_frame)
                                interp_heading = (1 - alpha) * self.root_headings[
                                    left_frame
                                ] + alpha * self.root_headings[right_frame]
                                interpolated_headings.append(interp_heading)
                            elif left_frame is not None:
                                interpolated_headings.append(self.root_headings[left_frame])
                            elif right_frame is not None:
                                interpolated_headings.append(self.root_headings[right_frame])

                    if len(interpolated_headings) == len(t):
                        result["global_root_heading"] = torch.tensor(interpolated_headings, dtype=torch.float32)

            return result

    def set_keyframe_visibility(self, keyframe_idx: int, visible: bool):
        """Sets the visibility of scene elements at the given keyframe index."""
        if keyframe_idx not in self.scene_elements:
            return

        scene_elements = self.scene_elements[keyframe_idx]
        if "waypoint" in scene_elements:
            waypoint = scene_elements["waypoint"]
            # WaypointMesh has sphere, annulus, arrow_base, arrow_head components
            if hasattr(waypoint, "sphere"):
                waypoint.sphere.visible = visible
            if hasattr(waypoint, "annulus") and waypoint.annulus is not None:
                waypoint.annulus.visible = visible
            if hasattr(waypoint, "arrow_base") and waypoint.arrow_base is not None:
                waypoint.arrow_base.visible = visible
            if hasattr(waypoint, "arrow_head") and waypoint.arrow_head is not None:
                waypoint.arrow_head.visible = visible

        if "label" in scene_elements:
            label = scene_elements["label"]
            if hasattr(label, "visible"):
                label.visible = visible

    def set_interval_labels_visibility(self, frame_idx: int):
        # set interval labels visibility
        for interval in self.interval_labels.keys():
            interval_label = self.interval_labels[interval]
            start_frame_idx, end_frame_idx = interval
            visibility = frame_idx <= end_frame_idx
            interval_label.visible = visibility

    def clear(self, frame_idx: Optional[int] = None):
        frame_idx_list = list(self.keyframes.keys()) if frame_idx is None else [frame_idx]
        for fidx in frame_idx_list:
            self.scene_elements[fidx]["waypoint"].clear()
            if "label" in self.scene_elements[fidx]:
                self.server.scene.remove_by_name(self.scene_elements[fidx]["label"].name)

            self.keyframes.pop(fidx)
            self.scene_elements.pop(fidx)
            # Also clear heading if it exists
            if fidx in self.root_headings:
                self.root_headings.pop(fidx)

        if frame_idx is None:
            # clear all interval labels if clearing all keyframes
            for interval_label in list(self.interval_labels.values()):
                self.server.scene.remove_by_name(interval_label.name)
            self.interval_labels.clear()

            # clear line segments if turning off dense path
            if self.line_segments is not None:
                self.server.scene.remove_by_name(self.line_segments.name)
                self.line_segments = None

            # Clear all headings when clearing everything
            self.root_headings.clear()

            # Invalidate cache when clearing all keyframes
            self._cached_t = None
            self._cached_path3d = None
        else:
            # Invalidate cache when clearing a specific keyframe
            self._cached_t = None
            self._cached_path3d = None


def load_example_cases(examples_base_dir):
    example_dirs = os.listdir(examples_base_dir)
    example_names = sorted([dir for dir in example_dirs if os.path.isdir(os.path.join(examples_base_dir, dir))])
    example_dict = {name: os.path.join(examples_base_dir, name) for name in example_names}
    return example_dict