JackBot/ml/MetricsOverlay.py

"""
ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
"""
from typing import List, Tuple, Optional
import numpy as np
import pybullet as p


class MetricsHUD:
    """Renders real-time telemetry as black text floating in 3D, always facing the active camera."""

    def __init__(self, physics_client_id: int = 0):
        self.client_id = physics_client_id
        self._text_id: Optional[int] = None

    def _get_camera_facing_orientation(self) -> List[float]:
        """Calculates a quaternion that points the text towards the current GUI camera."""
        try:
            cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id)
            # cam_info index 8: yaw, index 9: pitch
            yaw = cam_info[8]
            pitch = cam_info[9]
            
            # Orient the text normal toward the camera view direction
            # PyBullet text default faces local +Z/-Y depending on roll,
            # converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians)
            roll_rad = 0.0
            pitch_rad = np.radians(pitch + 90.0)
            yaw_rad = np.radians(yaw)
            
            text_orientation = p.getQuaternionFromEuler(
                [pitch_rad, roll_rad, yaw_rad], 
                physicsClientId=self.client_id
            )
            return text_orientation
        except Exception:
            # Fallback default orientation if camera info call fails
            return [0.0, 0.0, 0.0, 1.0]

    def update(
        self,
        episode: int,
        step: int,
        robot_rewards: List[float],
        cmd_vel: np.ndarray,
        fps: float = 0.0,
        avg_height: float = 0.0,
        roll_pitch: Tuple[float, float] = (0.0, 0.0)
    ) -> None:
        """Updates floating black text block in 3D space with billboarding."""
        sorted_rewards = sorted(robot_rewards, reverse=True)
        top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
        top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00"
        top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00"

        vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
        vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
        omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0

        hud_text = (
            f"=== JACKBOT METRICS ===\n"
            f"Episode:     {episode}\n"
            f"Global Step: {step}\n"
            f"FPS:         {fps:.1f}\n"
            f"----------------------\n"
            f"Top Rewards: [{top1}, {top2}, {top3}]\n"
            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
            f"Height:      {avg_height:.3f} m\n"
            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
        )

        # Position above origin in simulation world
        text_position = [-0.8, -0.8, 1.2]
        text_color = [0, 0, 0]  # Pure black

        # Calculate dynamic orientation to align text flat against camera plane
        text_orientation = self._get_camera_facing_orientation()

        if self._text_id is None:
            self._text_id = p.addUserDebugText(
                text=hud_text,
                textPosition=text_position,
                textColorRGB=text_color,
                textSize=0.1,
                textOrientation=text_orientation,
                physicsClientId=self.client_id
            )
        else:
            self._text_id = p.addUserDebugText(
                text=hud_text,
                textPosition=text_position,
                textColorRGB=text_color,
                textSize=0.1,
                textOrientation=text_orientation,
                replaceItemUniqueId=self._text_id,
                physicsClientId=self.client_id
            )

    def reset(self) -> None:
        """Removes the active debug text item so a new episode starts with a clean overlay."""
        if self._text_id is not None:
            try:
                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
            except Exception:
                pass
        self._text_id = None


class LeaderCrown:
    """Renders a floating crown or star emoji above the leading robot in PyBullet."""

    def __init__(self, physics_client_id: int = 0):
        self.client_id = physics_client_id
        self._text_id = None

    def update(self, leader_pos: list[float]):
        """Positions a floating crown ~0.35m directly above the lead robot's base."""
        crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35]
        
        # You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1"
        crown_text = "👑"

        if self._text_id is None:
            self._text_id = p.addUserDebugText(
                text=crown_text,
                textPosition=crown_pos,
                textColorRGB=[1.0, 0.84, 0.0],
                textSize=2.0,
                physicsClientId=self.client_id
            )
        else:
            self._text_id = p.addUserDebugText(
                text=crown_text,
                textPosition=crown_pos,
                textColorRGB=[1.0, 0.84, 0.0],
                textSize=2.0,
                replaceItemUniqueId=self._text_id,
                physicsClientId=self.client_id
            )

    def reset(self):
        if self._text_id is not None:
            try:
                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
            except Exception:
                pass
        self._text_id = None