"""A tiny MuJoCo world: one bouncing ball constrained to vertical motion.

Run: python3 examples/mujoco_first_steps.py
The bouncing-ball workshop uses this same ball, contact, and timestep.
"""

import mujoco


TIME_STEP = 0.001   # seconds per simulation step
STEPS = 1500        # 1.5 s of simulated time
RADIUS = 0.05       # metres
CONTACT = 'solref="-100000 -20" solimp="0.99 0.99 0.001"'   # a stiff, lightly damped spring


def make_model(gravity=9.81, contact=CONTACT):
    xml = f"""
    <mujoco model="one bouncing ball">
      <option timestep="{TIME_STEP}" gravity="0 0 -{gravity}"/>
      <default>
        <geom {contact}/>
      </default>
      <worldbody>
        <geom name="floor" type="plane" size="2 2 0.1"/>
        <body name="ball">
          <joint name="up_down" type="slide" axis="0 0 1"/>
          <geom type="sphere" size="{RADIUS}" mass="0.1"/>
        </body>
      </worldbody>
    </mujoco>
    """
    return mujoco.MjModel.from_xml_string(xml)


def run_drop(gravity=9.81, start_height=1.0, contact=CONTACT):
    model = make_model(gravity, contact)  # The world's rules: gravity, shapes, contact, timestep.
    data = mujoco.MjData(model)           # One changing state of that world.
    data.qpos[0] = start_height
    data.qvel[0] = 0.0
    mujoco.mj_forward(model, data)        # Recalculate from this start; time does not advance.

    rows = [(data.time, data.qpos[0], data.qvel[0], data.ncon > 0)]
    for _ in range(STEPS):
        mujoco.mj_step(model, data)       # Advance time and update data in place.
        mujoco.mj_forward(model, data)    # Refresh contacts for the new position.
        rows.append((data.time, data.qpos[0], data.qvel[0], data.ncon > 0))
    return rows


if __name__ == "__main__":
    rows = run_drop()
    for step in (0, 1, 100, 400, 430, 445, 460, 1500):
        time, height, velocity, contact = rows[step]
        print(f"step {step:4d}: time {time:.3f} s  height {height:.5f} m  "
              f"velocity {velocity:+.4f} m/s  contact {contact}")