"""Convert the arm URDF -> a standalone, viewable MuJoCo MJCF (adds a floor, light, checker ground).
Usage: make_mjcf.py <in.urdf> <out.xml>   (STLs must sit alongside the xml). View with:
  python -m mujoco.viewer --mjcf=out.xml   (or drop it in the `simulate` app)."""
import sys
import xml.etree.ElementTree as ET
import mujoco

urdf = sys.argv[1] if len(sys.argv) > 1 else "two_servo.urdf"
out = sys.argv[2] if len(sys.argv) > 2 else "two_servo.xml"

# 1) load URDF into MuJoCo (needs inertials + a mujoco/compiler block)
root = ET.parse(urdf).getroot()
# face the same way as the YAM arm: +180 deg about vertical (Z), on top of the Y-up->Z-up (Rx90) root rot
for j in root.findall("joint"):
    if j.get("name") == "world_to_servo_0":
        j.find("origin").set("rpy", "1.5708 0 3.14159")
# the driven revolute joints -> position actuators (control sliders + hold against gravity)
jlims = [(j.get("name"), j.find("limit").get("lower"), j.find("limit").get("upper"))
         for j in root.findall("joint") if j.get("type") == "revolute"]
for l in root.findall("link"):
    if l.find("inertial") is None:
        it = ET.SubElement(l, "inertial")
        ET.SubElement(it, "mass", {"value": "0.05"})
        ET.SubElement(it, "inertia", dict(ixx="1e-5", ixy="0", ixz="0", iyy="1e-5", iyz="0", izz="1e-5"))
mj = ET.Element("mujoco")
cc = ET.SubElement(mj, "compiler")
cc.set("discardvisual", "false")
cc.set("meshdir", ".")
root.insert(0, mj)
model = mujoco.MjModel.from_xml_string(ET.tostring(root, encoding="unicode"))
mujoco.mj_saveLastXML(out, model)

# 2) augment the saved MJCF for nice viewing: floor + light + checker + framing
tree = ET.parse(out)
r = tree.getroot()
r.set("model", "smith300_arm")
asset = r.find("asset")
if asset is None:
    asset = ET.SubElement(r, "asset")
ET.SubElement(asset, "texture", {"name": "sky", "type": "skybox", "builtin": "gradient",
                                 "rgb1": ".3 .5 .7", "rgb2": "0 0 0", "width": "512", "height": "512"})
ET.SubElement(asset, "texture", {"name": "grid", "type": "2d", "builtin": "checker", "width": "512",
                                 "height": "512", "rgb1": ".1 .15 .2", "rgb2": ".2 .25 .3"})
ET.SubElement(asset, "material", {"name": "grid", "texture": "grid", "texrepeat": "6 6", "reflectance": ".2"})
wb = r.find("worldbody")
ET.SubElement(wb, "light", {"pos": "0.3 -0.3 1.2", "dir": "-0.3 0.3 -1", "directional": "true"})
ET.SubElement(wb, "geom", {"name": "floor", "type": "plane", "size": "1.5 1.5 .05", "material": "grid",
                           "pos": "0 0 0", "contype": "0", "conaffinity": "0"})
ET.SubElement(r, "statistic", {"center": "-0.15 0.2 0.2", "extent": "0.7"})
vis = ET.SubElement(r, "visual")
ET.SubElement(vis, "global", {"azimuth": "130", "elevation": "-20", "offwidth": "1600", "offheight": "1600"})
# position actuators: these show as control sliders in the viewer AND hold each joint against gravity
act = ET.SubElement(r, "actuator")
for name, lo, hi in jlims:
    ET.SubElement(act, "position", {"name": "drive_" + name, "joint": name,
                                    "kp": "40", "kv": "2", "ctrlrange": f"{lo} {hi}"})
tree.write(out)

# 3) verify it reloads
m2 = mujoco.MjModel.from_xml_path(out)
print(f"wrote {out}: nbody={m2.nbody} njnt={m2.njnt} ngeom={m2.ngeom} nmesh={m2.nmesh}")