#!/usr/bin/env python3
"""Viewer test: load the Panda MuJoCo model with exoskeleton ArUco board and display it.

No rosbridge connection needed -- just loads the model, positions the exo meshes
at the robot's FK state, and launches an interactive MuJoCo viewer.

Usage:
  python hand_eye_calib/viewer_test.py
"""
import sys
import os
import signal

# Add parent directory so we can import ExoConfigs and exo_utils
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

import mujoco
import numpy as np

from panda_exo_handeye import PANDA_HANDEYE_CONFIG
from exo_utils import position_exoskeleton_meshes, get_link_poses_from_robot

N_ARM_JOINTS = 7
GRIPPER_POS_MAX = 0.04  # finger joint range 0-0.04 m per finger


def main() -> int:
    robot_config = PANDA_HANDEYE_CONFIG

    print(f"Loading model: {robot_config.base_xml_path}")
    model = mujoco.MjModel.from_xml_string(robot_config.xml)
    data = mujoco.MjData(model)

    # Set a non-zero home pose so the arm isn't folded at origin
    # Standard Panda home: slightly bent elbow, hand forward
    home_qpos = np.array([0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785])
    n_arm = min(N_ARM_JOINTS, data.qpos.size)
    data.qpos[:n_arm] = home_qpos[:n_arm]
    data.ctrl[:n_arm] = home_qpos[:n_arm]

    # Open gripper
    has_gripper = data.qpos.size >= N_ARM_JOINTS + 2
    if has_gripper:
        data.qpos[N_ARM_JOINTS] = GRIPPER_POS_MAX
        data.qpos[N_ARM_JOINTS + 1] = GRIPPER_POS_MAX

    mujoco.mj_forward(model, data)

    # Position exoskeleton meshes at FK-derived link poses
    link_poses = get_link_poses_from_robot(robot_config, model, data)
    position_exoskeleton_meshes(robot_config, model, data, link_poses)

    print("Launching MuJoCo viewer (close window to exit)...")
    viewer = mujoco.viewer.launch_passive(model, data)

    def shutdown(*_):
        viewer.close()

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    while viewer.is_running():
        mujoco.mj_step(model, data)
        position_exoskeleton_meshes(
            robot_config, model, data, get_link_poses_from_robot(robot_config, model, data)
        )
        viewer.sync()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())