"""
ball.py — the bouncing-ball world. Everything here is fixed physics; nothing is learned.

It is the ball from the MuJoCo basics lesson (radius 5 cm, 100 g, the same contact spring, 1 ms steps)
with one addition: a sideways joint, so a throw can travel across the floor.

State of the ball, used by every other script:
    s = [x, z, vx, vz]   (m, m, m/s, m/s)   vz > 0 means moving up
The velocities are MuJoCo's own state (data.qvel), not estimated from positions.
"""

import numpy as np
import mujoco

SIM_DT = 0.001                  # one MuJoCo step: 1 ms (the bounce needs small steps)
SUBSTEPS = 10                   # we keep every 10th state ...
MODEL_DT = SIM_DT * SUBSTEPS    # ... so the network sees the ball every 10 ms
RADIUS = 0.05
CONTACT = 'solref="-100000 -20" solimp="0.99 0.99 0.001"'   # a stiff, lightly damped spring

SCENE_XML = f"""
<mujoco model="bouncing ball">
  <option timestep="{SIM_DT}" gravity="0 0 -9.81"/>
  <default>
    <geom friction="0 0 0" {CONTACT}/>
  </default>
  <worldbody>
    <geom name="floor" type="plane" size="50 5 0.1"/>
    <body name="ball">
      <joint name="x" type="slide" axis="1 0 0"/>
      <joint name="z" type="slide" axis="0 0 1"/>
      <geom type="sphere" size="{RADIUS}" mass="0.1"/>
    </body>
  </worldbody>
</mujoco>
"""


def simulate(s0, n_steps):
    """Throw the ball from state s0 and record n_steps states, 10 ms apart.

    Returns states (n_steps + 1, 4) and contact (n_steps,): True if the ball touched the floor
    at any moment between two recorded states.
    """
    model = mujoco.MjModel.from_xml_string(SCENE_XML)
    data = mujoco.MjData(model)
    data.qpos[:] = s0[:2]           # x, z
    data.qvel[:] = s0[2:]           # vx, vz
    mujoco.mj_forward(model, data)

    states = np.zeros((n_steps + 1, 4))
    contact = np.zeros(n_steps, dtype=bool)
    states[0] = s0
    for t in range(n_steps):
        for _ in range(SUBSTEPS):
            mujoco.mj_step(model, data)
            mujoco.mj_forward(model, data)      # refresh contacts for the new position
            contact[t] |= data.ncon > 0
        states[t + 1] = [*data.qpos, *data.qvel]
    return states, contact