"""Render the gripper CLOSED (jaw=0) vs OPEN (jaw=OPEN_ANGLE) side by side, to check kinematics + flush."""
import sys
import xml.etree.ElementTree as ET
import mujoco
import numpy as np
from PIL import Image

urdf = sys.argv[1] if len(sys.argv) > 1 else "gripper_testing.urdf"
out = sys.argv[2] if len(sys.argv) > 2 else "gripper_test.png"
angles = [float(a) for a in sys.argv[3].split(",")] if len(sys.argv) > 3 else [0.0, 0.9]

root = ET.parse(urdf).getroot()
for l in root.findall("link"):
    if l.find("inertial") is None:
        it = ET.SubElement(l, "inertial")
        ET.SubElement(it, "mass", {"value": "0.1"})
        ET.SubElement(it, "inertia", dict(ixx="1e-4", ixy="0", ixz="0", iyy="1e-4", iyz="0", izz="1e-4"))
mj = ET.Element("mujoco")
cc = ET.SubElement(mj, "compiler")
cc.set("discardvisual", "false")
cc.set("meshdir", ".")
root.insert(0, mj)
m = mujoco.MjModel.from_xml_string(ET.tostring(root, encoding="unicode"))
d = mujoco.MjData(m)
adr = m.jnt_qposadr[mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, "jaw")]

views = [(90, -10), (0, -10), (35, -25)]     # (azimuth, elevation)
rows = []
for az, el in views:
    ims = []
    for ang in angles:
        d.qpos[:] = 0
        d.qpos[adr] = ang
        mujoco.mj_forward(m, d)
        cam = mujoco.MjvCamera()
        cam.lookat[:] = [-0.055, 0.0, 0.0]
        cam.distance = 0.20
        cam.azimuth = az
        cam.elevation = el
        with mujoco.Renderer(m, 440, 440) as r:
            r.update_scene(d, cam)
            ims.append(r.render())
    rows.append(np.hstack(ims))
Image.fromarray(np.vstack(rows)).save(out)
print("wrote", out, "angles", angles)