#!/usr/bin/env python3
"""Generate symmetric box-geom stairs XML for MuJoCo sim2sim.

The course is a flat approach, ascending stairs, a top platform, and matching
descending stairs. Each tread is a solid box rising from the ground.

Usage:
    uv run scripts/gen_stairs_box.py                           # default: 10 steps × 6cm = 60cm
    uv run scripts/gen_stairs_box.py --step-height 0.04 --num-steps 5
    uv run scripts/gen_stairs_box.py --step-height 0.10 --num-steps 8 --step-depth 0.4
"""
import argparse, os

TPL = '''<mujoco model="go1 box stairs scene">
  <include file="go1_motor_actuator.xml" />
  <include file="materials.xml" />
  <statistic center="0 0 0.3" extent="2" meansize="0.04" />

  <visual>
    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
    <rgba haze="0.15 0.25 0.35 1" />
    <global azimuth="120" elevation="-20" />
    <map force="0.01" />
    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
    <quality shadowsize="8192" />
  </visual>

  <worldbody>
    <light pos="0 0 4" dir="0 0 -1" directional="true" />

    <geom name="floor" pos="0 0 -0.001" size="0 0 0.001" type="plane"
          material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />

{steps}
  </worldbody>
</mujoco>
'''

STEP_TPL = '    <geom name="{name}" type="box" size="{sx:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
PLAT_TPL = '    <geom name="platform" type="box" size="{sx:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--step-height", type=float, default=0.06, help="rise per step [m]")
    p.add_argument("--step-depth", type=float, default=0.30, help="tread depth per step [m]")
    p.add_argument("--num-steps", type=int, default=10, help="number of steps")
    p.add_argument("--platform-depth", type=float, default=1.0, help="top platform depth [m]")
    args = p.parse_args()

    h = args.step_height
    d = args.step_depth
    n = args.num_steps
    platform_depth = args.platform_depth

    steps_xml = ""
    for i in range(n):
        top = (i + 1) * h
        steps_xml += STEP_TPL.format(
            name=f"step_up_{i}", sx=d / 2, sz=top / 2,
            x=(i + 0.5) * d, z=top / 2,
        )

    # Platform at top
    total_height = n * h
    platform_start = n * d
    platform_end = platform_start + platform_depth
    steps_xml += PLAT_TPL.format(
        sx=platform_depth / 2, sz=total_height / 2,
        x=platform_start + platform_depth / 2, z=total_height / 2,
    )

    for i in range(n):
        top = (n - i - 1) * h
        if top <= 0:
            continue
        steps_xml += STEP_TPL.format(
            name=f"step_down_{i}", sx=d / 2, sz=top / 2,
            x=platform_end + (i + 0.5) * d, z=top / 2,
        )

    out = TPL.format(steps=steps_xml.rstrip())

    out_dir = os.path.join(os.path.dirname(__file__), "..",
                           "motrix_envs", "src", "motrix_envs", "locomotion",
                           "go1", "xmls")
    out_path = os.path.join(out_dir, "scene_stairs_box.xml")
    with open(out_path, "w") as f:
        f.write(out)

    max_h = n * h
    finish_x = platform_end + n * d
    print(f"Generated {n} steps up/down x {h*100:.0f}cm = {max_h*100:.0f}cm total")
    print(f"  step depth: {d*100:.0f}cm  platform: {platform_depth:.2f}m")
    print(f"  course: x=0.00m to x={finish_x:.2f}m")
    print(f"  saved: {out_path}")


if __name__ == "__main__":
    main()