"""Single clean hero render of an arm URDF (larger via offscreen-buffer override).
Usage: render_hero.py <urdf> <out.png> [az] [el] [q]
"""
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]
az = float(sys.argv[3]) if len(sys.argv) > 3 else 40.0
el = float(sys.argv[4]) if len(sys.argv) > 4 else -12.0
q = float(sys.argv[5]) if len(sys.argv) > 5 else 0.0
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)
if m.nq:
    d.qpos[:] = q
mujoco.mj_forward(m, d)
cam = mujoco.MjvCamera()
cam.lookat[:] = m.stat.center
cam.distance = m.stat.extent * 1.45
cam.azimuth = az
cam.elevation = el
with mujoco.Renderer(m, 470, 630) as r:
    r.update_scene(d, cam)
    Image.fromarray(r.render()).save(out)
print("rendered", out, "az", az, "el", el)