From 97ed5692a45ce07146fdd3b8e03fa92ae5129982 Mon Sep 17 00:00:00 2001
From: budzianowski <budzianowski@gmail.com>
Date: Thu, 25 Apr 2024 23:01:14 -0700
Subject: [PATCH] update scripts

---
 sim/scripts/create_fixed_urdf.py | 66 ++++++++++++++++++++++++++++++--
 1 file changed, 63 insertions(+), 3 deletions(-)

diff --git a/sim/scripts/create_fixed_urdf.py b/sim/scripts/create_fixed_urdf.py
index 653bc0dd..74da8a27 100644
--- a/sim/scripts/create_fixed_urdf.py
+++ b/sim/scripts/create_fixed_urdf.py
@@ -4,11 +4,13 @@
 import xml.etree.ElementTree as ET
 from sim.stompy.joints import StompyFixed
 
+STOMPY_URDF = "stompy/robot.urdf"
+STOMPY_MJCF = "stompy/robot.xml"
+
 
 def update_urdf():
-    tree = ET.parse("stompy/robot.urdf")
+    tree = ET.parse()
     root = tree.getroot()
-
     stompy = StompyFixed()
 
     revolute_joints = set(stompy.default_standing().keys())
@@ -32,5 +34,63 @@ def update_urdf():
     tree.write("stompy/robot_fixed.urdf")
 
 
+def update_mjcf():
+    tree = ET.parse(STOMPY_MJCF)
+    root = tree.getroot()
+    # Create light element
+    light = ET.Element("light", {
+        "cutoff": "100",
+        "diffuse": "1 1 1",
+        "dir": "-0 0 -1.3",
+        "directional": "true",
+        "exponent": "1",
+        "pos": "0 0 1.3",
+        "specular": ".1 .1 .1"
+    })
+
+    # Create floor geometry element
+    floor = ET.Element("geom", {
+        "conaffinity": "1",
+        "condim": "3",
+        "name": "floor",
+        "pos": "0 0 -1.5",
+        "rgba": "0.8 0.9 0.8 1",
+        "size": "40 40 40",
+        "type": "plane",
+        "material": "MatPlane"
+    })
+    # Access the <worldbody> element
+    worldbody = root.find("worldbody")
+    # Insert new elements at the beginning of the worldbody
+    worldbody.insert(0, floor)
+    worldbody.insert(0, light)
+
+    new_root_body = ET.Element("body", name="root", pos="-1.4 0 -0.3", quat="1 0 0 1")
+
+    # List to store items to be moved to the new root body
+    items_to_move = []
+    # Gather all children (geoms and bodies) that need to be moved under the new root body
+    for element in worldbody:
+        items_to_move.append(element)
+    # Move gathered elements to the new root body
+    for item in items_to_move:
+        worldbody.remove(item)
+        new_root_body.append(item)
+    # Add the new root body to the worldbody
+    worldbody.append(new_root_body)
+    worldbody2 = worldbody
+    # # Create a new root worldbody and add the new root body to it
+    # new_worldbody = ET.Element("worldbody", name="new_root")
+    index = list(root).index(worldbody)
+
+    # # Remove the old <worldbody> and insert the new one at the same index
+    root.remove(worldbody)
+    root.insert(index, worldbody2)
+    modified_xml = ET.tostring(root, encoding="unicode")
+
+    with open("stompy/robot_fixed.xml", "w") as f:
+        f.write(modified_xml)
+
+
 if __name__ == "__main__":
-    update_urdf()
+    # update_urdf()