"""Servo TRANSLUCENT + holder SOLID, so any holder poke-through into the servo is
directly visible. Usage: render_xray.py <servo.stl> <holder.stl> <out.png>
"""
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"/><light pos="0.3 0.2 0.1"/>
    <geom type="mesh" mesh="servo" rgba="0.55 0.55 0.60 0.30"/>
    <geom type="mesh" mesh="holder" rgba="0.20 0.55 0.92 1.0"/>
  </worldbody>
</mujoco>"""
m = mujoco.MjModel.from_xml_string(xml)
d = mujoco.MjData(m)
mujoco.mj_forward(m, d)
ims = []
# overview angles + close-ups on the +X flange / bolt-boss regions
shots = [
    (35, -18, (-0.01, -0.009, 0.0), 0.13),
    (90, -10, (-0.005, -0.009, 0.0), 0.115),
    (90, -88, (-0.005, -0.009, 0.0), 0.115),
    (20, 8, (0.004, 0.0, 0.010), 0.055),     # close-up: +Y flange screw bosses
    (20, 8, (0.006, -0.018, -0.010), 0.055),  # close-up: -Y flange screw bosses
]
for az, el, la, dist in shots:
    cam = mujoco.MjvCamera()
    cam.lookat[:] = la
    cam.distance = dist
    cam.azimuth = az
    cam.elevation = el
    with mujoco.Renderer(m, 460, 460) as r:
        r.update_scene(d, cam)
        ims.append(r.render())
Image.fromarray(np.concatenate(ims, 1)).save(out)
print("rendered", out)