"""Render a URDF at a SPECIFIC per-joint pose. Usage:
   so100_render.py <urdf> <out.png> "joint=angle,joint2=angle2,..."   (unset joints stay at 0)
"""
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]
poses = {}
if len(sys.argv) > 3 and sys.argv[3]:
    for kv in sys.argv[3].split(","):
        k, v = kv.split("=")
        poses[k.strip()] = float(v)

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")
c = ET.SubElement(mj, "compiler")
c.set("discardvisual", "false")
c.set("meshdir", ".")
root.insert(0, mj)
m = mujoco.MjModel.from_xml_string(ET.tostring(root, encoding="unicode"))
d = mujoco.MjData(m)
for name, ang in poses.items():
    try:
        d.qpos[m.joint(name).qposadr[0]] = ang
    except KeyError:
        print("WARN: no joint", name)
mujoco.mj_forward(m, d)
print("pose", poses)
# optional 4th arg: focus a link name -> frame tightly on it (e.g. the gripper). 5th: distance (m).
focus = sys.argv[4] if len(sys.argv) > 4 else None
fdist = float(sys.argv[5]) if len(sys.argv) > 5 else 0.13
lookat = d.xpos[m.body(focus).id].copy() if focus else m.stat.center
dist = fdist if focus else m.stat.extent * 1.5
ims = []
for az, el in [(45, -20), (135, -18), (90, -89), (0, -5)]:
    cam = mujoco.MjvCamera()
    cam.lookat[:] = lookat
    cam.distance = dist
    cam.azimuth, cam.elevation = az, el
    with mujoco.Renderer(m, 470, 470) as r:
        r.update_scene(d, cam)
        ims.append(r.render())
Image.fromarray(np.concatenate(ims, 1)).save(out)
print("rendered", out)