"""Render a servo STL + holder STL overlaid (holder translucent) from a few angles.
Usage: render_overlay.py <servo.stl> <holder.stl> <out.png>
Meshes assumed in the STEP frame (mm); scaled to meters for MuJoCo.
"""
import sys
import mujoco
import numpy as np
from PIL import Image

servo, holder, out = sys.argv[1], sys.argv[2], sys.argv[3]
xml = f"""<mujoco>
  <compiler meshdir="." />
  <asset>
    <mesh name="servo" file="{servo}" scale="0.001 0.001 0.001"/>
    <mesh name="holder" file="{holder}" scale="0.001 0.001 0.001"/>
  </asset>
  <worldbody>
    <light pos="0 0 0.5"/>
    <geom type="mesh" mesh="servo" rgba="0.32 0.32 0.36 1"/>
    <geom type="mesh" mesh="holder" rgba="0.20 0.62 0.92 0.45"/>
  </worldbody>
</mujoco>"""
m = mujoco.MjModel.from_xml_string(xml)
d = mujoco.MjData(m)
mujoco.mj_forward(m, d)
ims = []
for az, el in [(45, -20), (135, -18), (90, -88), (0, 0)]:
    cam = mujoco.MjvCamera()
    cam.lookat[:] = (-0.013, -0.009, 0.0)
    cam.distance = 0.115
    cam.azimuth = az
    cam.elevation = el
    with mujoco.Renderer(m, 480, 480) as r:
        r.update_scene(d, cam)
        ims.append(r.render())
Image.fromarray(np.concatenate(ims, 1)).save(out)
print("rendered", out)