"""Open the arm, wall and target in MuJoCo's interactive viewer window.

Run: python3 classes/ml/labs/robot-arm/view_arm.py

In the window:
- drag with the left mouse button to turn the camera, right button to move it, scroll to zoom;
- double-click a link to select it, then hold Ctrl and drag with the right button to push it into the wall;
- open "Control" in the right-hand panel: its two sliders are the motors' target angles (radians), as the
  four commands set them. Move one and watch the arm chase it.
Physics runs live; red marks show where things touch.
"""
import math
import time

import mujoco
import mujoco.viewer
import numpy as np

import mujoco_arm as M
import world as W


def lit_xml():
    """The lesson's scene with lights, colours, a side camera and the target as a ball (the lesson's has none)."""
    xml = M.XML.replace('<option timestep="0.001" gravity="0 0 -9.81"/>', """<option timestep="0.001" gravity="0 0 -9.81"/>
  <visual><global offwidth="640" offheight="360"/><headlight ambient="0.35 0.35 0.35" diffuse="0.5 0.5 0.5"/>
    <scale contactwidth="0.02" contactheight="0.006"/><rgba contactpoint="0.9 0.15 0.1 1"/></visual>
  <asset><texture name="grid" type="2d" builtin="checker" rgb1="0.93 0.93 0.9" rgb2="0.84 0.84 0.8" width="200" height="200"/>
    <material name="floor" texture="grid" texrepeat="20 12"/>
    <texture name="sky" type="skybox" builtin="gradient" rgb1="0.97 0.97 0.95" rgb2="0.8 0.84 0.88" width="64" height="64"/></asset>""")
    xml = xml.replace('<geom name="floor" type="plane" size="0.3 0.1 0.01" pos="0 0 -0.006"/>',
                      '<geom name="floor" type="plane" size="0.6 0.4 0.01" pos="0 0 -0.006" material="floor"/>'
                      '<light pos="0.05 -0.3 0.4" dir="0 0.6 -1" diffuse="0.6 0.6 0.6"/>'
                      '<camera name="side" pos="0.06 -0.3 0.08" xyaxes="1 0 0 0 0.1 1"/>'
                      f'<geom type="sphere" size="0.015" pos="{W.TARGET[0] / 100} 0 {W.TARGET[1] / 100}" '
                      'rgba="0.95 0.45 0.1 0.45" contype="0" conaffinity="0"/>')
    xml = xml.replace('<geom name="wall" type="box"', '<geom name="wall" rgba="0.45 0.47 0.5 1" type="box"')
    return xml.replace('size="0.004" mass', 'rgba="0.18 0.45 0.85 1" size="0.004" mass')


if __name__ == "__main__":
    model = mujoco.MjModel.from_xml_string(lit_xml())
    model.actuator_ctrllimited[:] = 1                                  # slider ranges: the joints' 0…180° and −180…0°
    model.actuator_ctrlrange[:] = [[0, math.pi], [-math.pi, 0]]
    data = mujoco.MjData(model)
    data.qpos[:] = data.ctrl[:] = np.radians(M.pose(W.HOME))           # start at home, as in the lessons
    with mujoco.viewer.launch_passive(model, data) as viewer:
        viewer.cam.lookat[:], viewer.cam.distance = [0.05, 0, 0.05], 0.6   # look at the arm from the side
        viewer.cam.azimuth, viewer.cam.elevation = 90, -10
        while viewer.is_running():                                      # our own loop: physics at real-time speed
            start = time.time()
            mujoco.mj_step(model, data)
            viewer.sync()                                               # show the new state, read the mouse and sliders
            time.sleep(max(0.0, model.opt.timestep - (time.time() - start)))
    time.sleep(0.5)                                                     # let the window thread finish