"""Overlay servo + holder + yoke to check collisions/fit.
Usage: render_three.py <servo.stl> <holder.stl> <yoke.stl> <out.png>
All in the STEP frame (mm); scaled to meters for MuJoCo.
"""
import sys
import mujoco
import numpy as np
from PIL import Image

servo, holder, yoke, out = sys.argv[1:5]
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"/>
    <mesh name="yoke" file="{yoke}" 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.40"/>
    <geom type="mesh" mesh="yoke" rgba="0.95 0.6 0.15 0.55"/>
  </worldbody>
</mujoco>"""
m = mujoco.MjModel.from_xml_string(xml)
d = mujoco.MjData(m)
mujoco.mj_forward(m, d)
ims = []
for az, el in [(35, -18), (90, -10), (90, -88), (-25, 5)]:
    cam = mujoco.MjvCamera()
    cam.lookat[:] = (-0.018, -0.009, 0.0)
    cam.distance = 0.135
    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)