"""
ml/MetricsOverlay.py - 3D HUD overlay for live simulation telemetry.

This file provides a small PyBullet HUD renderer that draws live robot metrics in the
simulation scene. It is used during GUI runs to show the current phase, command vector,
reward information, and basic motion statistics without leaving the 3D view.
"""
from typing import List, Tuple, Optional, Dict
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]
            
            pitch_rad = np.radians(pitch + 90.0)
            yaw_rad = np.radians(yaw)
            
            text_orientation = p.getQuaternionFromEuler(
                [pitch_rad, 0.0, yaw_rad], 
                physicsClientId=self.client_id
            )
            return text_orientation
        except Exception:
            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,
        height: float = 0.0,
        roll_pitch: Tuple[float, float] = (0.0, 0.0),
        mode: str = "direct",
        phase: str = "STAND_ONLY",
        distance: float = 0.0,
        status: str = "ALIVE",
        reward_components: Optional[Dict[str, float]] = None,
        ep_step: int = 0,
    ) -> None:
        """Updates floating text block in 3D space with expanded telemetry."""
        sorted_rewards = sorted(robot_rewards, reverse=True)
        top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 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[2] if len(cmd_vel) > 2 else 0.0

        lines = [
            "=== JACKBOT TELEMETRY ===",
            f"Mode:        {mode.upper()}",
            f"Curriculum:  {phase}",
            f"Status:      {status}",
            f"Episode:     {episode} (Step {ep_step})",
            f"Global Step: {step}",
            f"FPS:         {fps:.1f}",
            "-------------------------",
            f"Episode Rew: {top1}",
            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
            f"Height:      {height:.3f} m",
            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
            f"Max Dist:    {distance:.2f} m",
        ]

        if reward_components:
            lin_v = reward_components.get("lin_vel", 0.0)
            stab = reward_components.get("stability", 0.0)
            h_rew = reward_components.get("height", 0.0)
            jit = reward_components.get("jitter_penalty", 0.0)
            lines.append("--- Reward Components ---")
            lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
            lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")

        hud_text = "\n".join(lines)

        # 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()

        # Safely remove old text to prevent PyBullet ghosting/overlapping
        if self._text_id is not None:
            try:
                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
            except Exception:
                pass
                
        # Draw fresh text
        self._text_id = p.addUserDebugText(
            text=hud_text,
            textPosition=text_position,
            textColorRGB=text_color,
            textSize=0.085,
            textOrientation=text_orientation,
            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