"""Render the gripper open vs closed (arm at rest), zoomed on it, to verify actuation.
Usage: render_gripper.py <urdf> <out.png>
"""
import sys
import xml.etree.ElementTree as ET
import mujoco
import numpy as np
from PIL import Image

urdf, out = sys.argv[1], sys.argv[2]
root = ET.parse(urdf).getroot()
for l in root.findall("link"):
    if l.find("inertial") is None:
        i = ET.SubElement(l, "inertial")
        ET.SubElement(i, "mass", {"value": "0.1"})
        ET.SubElement(i, "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)
print("nq", m.nq)
gid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "gripper")
fidx = [m.jnt_qposadr[mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, j)] for j in ("finger_l", "finger_r")]
ims = []
for fq in (0.0, 0.012):
    d.qpos[:] = 0
    for a in fidx:
        d.qpos[a] = fq
    mujoco.mj_forward(m, d)
    cam = mujoco.MjvCamera()
    cam.lookat[:] = d.xpos[gid]
    cam.distance = 0.09
    cam.azimuth = 90
    cam.elevation = -12
    with mujoco.Renderer(m, 440, 440) as r:
        r.update_scene(d, cam)
        ims.append(r.render())
Image.fromarray(np.concatenate(ims, 1)).save(out)
print("rendered", out, "(open | closed)")