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

Each step is a separate box with vertical rises — much steeper than hfield.

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}

    <!-- Fill under stairs -->
    <geom name="fill" type="box" size="{fill_sx} 10 {fill_sz}" pos="{fill_x} 0 {fill_z}" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
  </worldbody>
</mujoco>
'''

STEP_TPL = '    <geom name="step{n}" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" 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} 10 {sz}" pos="{x} 0 {z}" 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("--box-thickness", type=float, default=0.03, help="box half-height [m]")
    args = p.parse_args()

    h = args.step_height
    d = args.step_depth
    n = args.num_steps
    sz = args.box_thickness  # half-height of each box

    steps_xml = ""
    for i in range(n):
        x = i * d
        z = i * h + sz  # center of box = step top surface - sz
        steps_xml += STEP_TPL.format(n=i, sx=d/2, sz=sz, x=x, z=z)

    # Platform at top
    plat_x = n * d + 0.5
    plat_z = n * h + sz
    steps_xml += PLAT_TPL.format(sx=0.5, sz=sz, x=plat_x, z=plat_z)

    # Fill box under stairs
    total_depth = n * d
    total_height = n * h
    fill_sx = total_depth / 2
    fill_sz = total_height / 2
    fill_x = total_depth / 2
    fill_z = -fill_sz

    out = TPL.format(steps=steps_xml.rstrip(),
                     fill_sx=fill_sx, fill_sz=fill_sz,
                     fill_x=fill_x, fill_z=fill_z)

    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
    print(f"Generated {n} steps × {h*100:.0f}cm = {max_h*100:.0f}cm total")
    print(f"  step depth: {d*100:.0f}cm  box thickness: {sz*200:.0f}cm")
    print(f"  saved: {out_path}")


if __name__ == "__main__":
    main()