"""MuJoCo native viewer debug visualizer implementation."""

from __future__ import annotations

import mujoco
import numpy as np
from typing_extensions import override

from mjlab.viewer.debug_visualizer import DebugVisualizer


class MujocoNativeDebugVisualizer(DebugVisualizer):
  """Debug visualizer for MuJoCo's native viewer.

  This implementation directly adds geometry to the MuJoCo scene using mjv_addGeoms
  and other MuJoCo visualization primitives.
  """

  def __init__(
    self,
    scn: mujoco.MjvScene,
    mj_model: mujoco.MjModel,
    env_idx: int,
    show_all_envs: bool = False,
  ):
    """Initialize the MuJoCo native visualizer.

    Args:
      scn: MuJoCo scene to add visualizations to
      mj_model: MuJoCo model for creating visualization data
      env_idx: Index of the environment being visualized
      show_all_envs: If True, visualize all environments instead of just env_idx
    """
    # DebugVisualizer interface.
    self.env_idx = env_idx
    self.show_all_envs = show_all_envs

    self._mjv_scene = scn
    self._initial_geom_count = scn.ngeom
    self._meansize: float = mj_model.stat.meansize

    # State used only by the ghost visualization. The model is a shared
    # reference and must not be modified; ghost transparency is controlled
    # through the mjVIS_TRANSPARENT flag rather than per-geom alpha.
    self._ref_mj_model = mj_model
    self._mj_data = mujoco.MjData(mj_model)
    self._vopt = mujoco.MjvOption()
    self._vopt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True
    self._pert = mujoco.MjvPerturb()

  @property
  @override
  def meansize(self) -> float:
    return self._meansize

  @override
  def add_arrow(
    self,
    start: np.ndarray,
    end: np.ndarray,
    color: tuple[float, float, float, float],
    width: float = 0.015,
    label: str | None = None,
  ) -> None:
    """Add an arrow visualization using MuJoCo's arrow geometry."""
    del label  # Unused.

    self._mjv_scene.ngeom += 1
    geom = self._mjv_scene.geoms[self._mjv_scene.ngeom - 1]
    geom.category = mujoco.mjtCatBit.mjCAT_DECOR

    mujoco.mjv_initGeom(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_ARROW.value,
      size=np.zeros(3),
      pos=np.zeros(3),
      mat=np.zeros(9),
      rgba=np.asarray(color, dtype=np.float32),
    )
    mujoco.mjv_connector(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_ARROW.value,
      width=width,
      from_=start,
      to=end,
    )

  @override
  def add_ghost_mesh(
    self,
    qpos: np.ndarray,
    model: mujoco.MjModel,
    mocap_pos: np.ndarray | None = None,
    mocap_quat: np.ndarray | None = None,
    alpha: float = 0.5,
    label: str | None = None,
  ) -> None:
    """Add a ghost mesh by rendering the robot at a different pose.

    This creates a semi-transparent copy of the robot geometry at the target pose.

    Args:
      qpos: Joint positions for the ghost pose
      model: MuJoCo model with pre-configured appearance (geom_rgba for colors).
        Must be structurally identical to the model this visualizer was
        created with, since it is evaluated against the internal MjData.
      mocap_pos: Optional mocap position(s) for fixed-base entities
      mocap_quat: Optional mocap quaternion(s) for fixed-base entities
      alpha: Transparency override (not used in MuJoCo implementation)
      label: Optional label (not used in MuJoCo implementation)
    """
    del alpha, label  # Unused.

    self._mj_data.qpos[:] = qpos
    if mocap_pos is not None and model.nmocap > 0:
      mocap_pos_arr = np.asarray(mocap_pos)
      if mocap_pos_arr.ndim == 1:
        self._mj_data.mocap_pos[0] = mocap_pos_arr
      else:
        self._mj_data.mocap_pos[:] = mocap_pos_arr
    if mocap_quat is not None and model.nmocap > 0:
      mocap_quat_arr = np.asarray(mocap_quat)
      if mocap_quat_arr.ndim == 1:
        self._mj_data.mocap_quat[0] = mocap_quat_arr
      else:
        self._mj_data.mocap_quat[:] = mocap_quat_arr
    mujoco.mj_forward(model, self._mj_data)

    mujoco.mjv_addGeoms(
      model,
      self._mj_data,
      self._vopt,
      self._pert,
      mujoco.mjtCatBit.mjCAT_DYNAMIC.value,
      self._mjv_scene,
    )

  @override
  def add_frame(
    self,
    position: np.ndarray,
    rotation_matrix: np.ndarray,
    scale: float = 0.3,
    label: str | None = None,
    axis_radius: float = 0.01,
    alpha: float = 1.0,
    axis_colors: tuple[tuple[float, float, float], ...] | None = None,
  ) -> None:
    """Add a coordinate frame visualization with RGB-colored axes."""
    del label  # Unused.

    default_colors = [(0.9, 0, 0), (0, 0.9, 0.0), (0.0, 0.0, 0.9)]
    colors = axis_colors if axis_colors is not None else default_colors

    for axis_idx in range(3):
      axis_direction = rotation_matrix[:, axis_idx]
      end_position = position + axis_direction * scale
      rgb = colors[axis_idx]
      color_rgba = (rgb[0], rgb[1], rgb[2], alpha)
      self.add_arrow(
        start=position,
        end=end_position,
        color=color_rgba,
        width=axis_radius,
      )

  @override
  def add_sphere(
    self,
    center: np.ndarray,
    radius: float,
    color: tuple[float, float, float, float],
    label: str | None = None,
  ) -> None:
    """Add a sphere visualization using MuJoCo's sphere geometry."""
    del label  # Unused.

    self._mjv_scene.ngeom += 1
    geom = self._mjv_scene.geoms[self._mjv_scene.ngeom - 1]
    geom.category = mujoco.mjtCatBit.mjCAT_DECOR

    mujoco.mjv_initGeom(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_SPHERE.value,
      size=np.array([radius, 0, 0]),
      pos=center,
      mat=np.eye(3).flatten(),
      rgba=np.asarray(color, dtype=np.float32),
    )

  @override
  def add_cylinder(
    self,
    start: np.ndarray,
    end: np.ndarray,
    radius: float,
    color: tuple[float, float, float, float],
    label: str | None = None,
  ) -> None:
    """Add a cylinder visualization using MuJoCo's cylinder connector."""
    del label  # Unused.

    self._mjv_scene.ngeom += 1
    geom = self._mjv_scene.geoms[self._mjv_scene.ngeom - 1]
    geom.category = mujoco.mjtCatBit.mjCAT_DECOR

    mujoco.mjv_initGeom(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_CYLINDER.value,
      size=np.zeros(3),
      pos=np.zeros(3),
      mat=np.zeros(9),
      rgba=np.asarray(color, dtype=np.float32),
    )
    mujoco.mjv_connector(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_CYLINDER.value,
      width=radius,
      from_=start,
      to=end,
    )

  @override
  def add_ellipsoid(
    self,
    center: np.ndarray,
    size: np.ndarray,
    mat: np.ndarray,
    color: tuple[float, float, float, float],
    label: str | None = None,
  ) -> None:
    """Add an ellipsoid visualization using MuJoCo's ellipsoid geometry."""
    del label  # Unused.

    self._mjv_scene.ngeom += 1
    geom = self._mjv_scene.geoms[self._mjv_scene.ngeom - 1]
    geom.category = mujoco.mjtCatBit.mjCAT_DECOR

    mujoco.mjv_initGeom(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_ELLIPSOID.value,
      size=np.asarray(size, dtype=np.float64),
      pos=np.asarray(center, dtype=np.float64),
      mat=np.asarray(mat, dtype=np.float64).flatten(),
      rgba=np.asarray(color, dtype=np.float32),
    )

  @override
  def add_box(
    self,
    center: np.ndarray,
    size: np.ndarray,
    mat: np.ndarray,
    color: tuple[float, float, float, float],
    label: str | None = None,
  ) -> None:
    """Add a box visualization using MuJoCo's box geometry."""
    del label  # Unused.

    self._mjv_scene.ngeom += 1
    geom = self._mjv_scene.geoms[self._mjv_scene.ngeom - 1]
    geom.category = mujoco.mjtCatBit.mjCAT_DECOR

    mujoco.mjv_initGeom(
      geom=geom,
      type=mujoco.mjtGeom.mjGEOM_BOX.value,
      size=np.asarray(size, dtype=np.float64),
      pos=np.asarray(center, dtype=np.float64),
      mat=np.asarray(mat, dtype=np.float64).flatten(),
      rgba=np.asarray(color, dtype=np.float32),
    )

  @override
  def clear(self) -> None:
    """Clear debug visualizations by resetting geom count."""
    self._mjv_scene.ngeom = self._initial_geom_count
