import mujoco, xml.etree.ElementTree as ET, numpy as np
from yam_load import yam_root

spec = mujoco.MjSpec()
spec.modelname = "arm_vs_yam"
spec.visual.global_.offwidth = 1600
spec.visual.global_.offheight = 1600
l = spec.worldbody.add_light(); l.pos = [0.3, -0.3, 2.0]; l.dir = [-0.3, 0.3, -1]
g = spec.worldbody.add_geom(); g.type = mujoco.mjtGeom.mjGEOM_PLANE; g.size = [2, 2, 0.05]; g.rgba = [.28, .3, .35, 1]
g.contype = 0; g.conaffinity = 0
ours = mujoco.MjSpec.from_file("two_servo.xml")
yam = mujoco.MjSpec.from_string(ET.tostring(yam_root(), encoding="unicode"))
f1 = spec.worldbody.add_frame(); f1.pos = [0, -0.35, 0]
spec.attach(ours, prefix="ours_", frame=f1)
f2 = spec.worldbody.add_frame(); f2.pos = [0, 0.35, 0]
spec.attach(yam, prefix="yam_", frame=f2)
model = spec.compile()
open("arm_vs_yam.xml", "w").write(spec.to_xml())
print("compiled OK: nbody", model.nbody, "njnt", model.njnt)

# render in-process (the saved xml needs bundled meshes; done in the export step)
from PIL import Image
d = mujoco.MjData(model)
for _ in range(800): mujoco.mj_step(model, d)
mujoco.mj_forward(model, d)
cam = mujoco.MjvCamera(); mujoco.mjv_defaultFreeCamera(model, cam)
cam.distance *= 1.15
with mujoco.Renderer(model, 1100, 750) as rr:
    rr.update_scene(d, cam); Image.fromarray(rr.render()).save("arm_vs_yam.png")
print("wrote arm_vs_yam.png")