import mujoco, xml.etree.ElementTree as ET, sys
Y = "/data/cameron/para/robot/yam/raiden_fork/third_party/exo_redo/robot_models/i2rt_yam_fidex"

def yam_root():
    t = ET.parse(f"{Y}/yam.xml"); r = t.getroot()
    c = r.find("compiler") or ET.SubElement(r, "compiler"); c.set("meshdir", f"{Y}/assets")
    v = r.find("visual") or ET.SubElement(r, "visual")
    g = v.find("global") or ET.SubElement(v, "global"); g.set("offwidth", "1600"); g.set("offheight", "1600")
    return r

if __name__ == "__main__":
    import numpy as np
    from PIL import Image
    r = yam_root()
    m = mujoco.MjModel.from_xml_string(ET.tostring(r, encoding="unicode")); d = mujoco.MjData(m); mujoco.mj_forward(m, d)
    print("YAM OK nbody", m.nbody, "nmesh", m.nmesh, "njnt", m.njnt)
    if len(sys.argv) > 1:
        cam = mujoco.MjvCamera(); mujoco.mjv_defaultFreeCamera(m, cam)
        with mujoco.Renderer(m, 620, 620) as rr:
            rr.update_scene(d, cam); Image.fromarray(rr.render()).save(sys.argv[1])
        print("wrote", sys.argv[1])