"""Robot basics: a two-link arm with position-controlled joints, simulated in MuJoCo.

Run: python3 site/ml/foundations/examples/robot_basics.py

The arm moves in a vertical plane, like the SO-101's shoulder_lift and elbow_flex joints. Each joint has a
motor that is given a target angle; its controller pushes toward that target. Angles on the page are in
degrees, as in a LeRobot SO-101 dataset; MuJoCo itself works in radians.
"""

import math

import mujoco
import numpy as np

L1, L2 = 0.10, 0.10          # link lengths in metres (the same as the robot-arm workshops)
TIME_STEP = 0.001            # seconds per physics step
FPS = 30                     # LeRobot records and commands at 30 frames per second

GAINS = {                    # kp: N·m per radian of error; kv: N·m per rad/s of speed
    "weak P + D": (1.5, 0.3),
    "strong P + D": (6.0, 0.4),
    "strong P, no D": (6.0, 0.0),
}


def make_model(kp, kv):
    xml = f"""
    <mujoco model="two-link arm">
      <option timestep="{TIME_STEP}" gravity="0 0 -9.81"/>
      <worldbody>
        <body name="upper_arm">
          <joint name="shoulder" type="hinge" axis="0 -1 0" armature="0.02"/>
          <geom type="capsule" fromto="0 0 0 {L1} 0 0" size="0.012" mass="0.20" contype="0" conaffinity="0"/>
          <body name="forearm" pos="{L1} 0 0">
            <joint name="elbow" type="hinge" axis="0 -1 0" armature="0.02"/>
            <geom type="capsule" fromto="0 0 0 {L2} 0 0" size="0.010" mass="0.15" contype="0" conaffinity="0"/>
            <site name="hand" pos="{L2} 0 0"/>
          </body>
        </body>
      </worldbody>
      <actuator>
        <position name="shoulder" joint="shoulder" kp="{kp}" kv="{kv}"/>
        <position name="elbow" joint="elbow" kp="{kp}" kv="{kv}"/>
      </actuator>
    </mujoco>
    """
    return mujoco.MjModel.from_xml_string(xml)


def angles(data):
    return np.degrees(data.qpos).copy()            # measured joint angles, in degrees


def push(data):
    """The shoulder motor's torque now: kp × (target − angle) − kv × speed, in N·m."""
    return data.actuator_force[0]


def step_response(kp, kv, start=(90.0, 0.0), target=(30.0, 0.0), seconds=2.0):
    """Hold the arm at `start`, then command `target` at time 0. Record the angles every physics step."""
    model = make_model(kp, kv)
    data = mujoco.MjData(model)
    data.qpos[:] = np.radians(start)
    data.ctrl[:] = np.radians(start)
    for _ in range(3000):                          # hold the start pose (straight up: gravity has no lever)
        mujoco.mj_step(model, data)
    data.qvel[:] = 0.0
    data.time = 0.0
    data.ctrl[:] = np.radians(target)              # the command: a target angle for each joint
    mujoco.mj_forward(model, data)                 # the motors react to the new targets
    rows = [(data.time, *angles(data), push(data))]
    while data.time < seconds - 1e-9:
        mujoco.mj_step(model, data)                # the position controller pushes toward the target
        mujoco.mj_forward(model, data)             # refresh the motor torque for the new angles
        rows.append((data.time, *angles(data), push(data)))
    return rows


def leader_angles(t):
    """A person moving the leader arm: reach out and up in 1 s, hold, come back in 1 s, hold."""
    home, reach = np.array([30.0, -100.0]), np.array([80.0, -40.0])

    def smooth(s):                                 # 0 → 1 with zero speed at both ends
        s = min(max(s, 0.0), 1.0)
        return 10 * s**3 - 15 * s**4 + 6 * s**5

    if t < 1.5:
        return home + (reach - home) * smooth((t - 0.5) / 1.0)
    return reach + (home - reach) * smooth((t - 2.0) / 1.0)


def teleoperate(kp, kv, seconds=3.5):
    """The LeRobot record loop at 30 Hz: read the follower, read the leader, send the leader's angles, save both."""
    model = make_model(kp, kv)
    data = mujoco.MjData(model)
    data.qpos[:] = np.radians(leader_angles(0.0))
    data.ctrl[:] = np.radians(leader_angles(0.0))
    for _ in range(3000):                          # settle at the start pose before recording
        mujoco.mj_step(model, data)
    data.qvel[:] = 0.0
    data.time = 0.0
    frames, trace = [], []
    for i in range(round(seconds * FPS)):
        t = i / FPS
        observation_state = angles(data)           # robot.get_observation(): the follower's measured angles
        action = leader_angles(t)                  # teleop.get_action(): the leader's angles
        data.ctrl[:] = np.radians(action)          # robot.send_action(): they become the follower's targets
        frames.append({"frame_index": i, "timestamp": t,
                       "observation.state": observation_state, "action": action})
        while data.time < (i + 1) / FPS - 1e-9:    # physics runs on between two frames
            mujoco.mj_step(model, data)
            trace.append((data.time, *angles(data)))
    return frames, trace


def raw_to_degrees(raw, range_min, range_max, resolution=4096):
    """LeRobot's degree reading for an STS3215 servo: 0° is the middle of the calibrated range."""
    mid = (range_min + range_max) / 2
    return (raw - mid) * 360 / (resolution - 1)


def summary(rows, target):
    """Start, furthest point, final angle and remaining gap of the shoulder, plus the time to cover 90% of the move."""
    t = np.array([r[0] for r in rows])
    shoulder = np.array([r[1] for r in rows])
    start, final = shoulder[0], shoulder[-1]
    moved = (shoulder - start) / (target - start)               # 0 at the start, 1 at the target
    furthest = shoulder[np.argmax(moved)]
    return {"start": start, "furthest": furthest, "final": final, "gap": target - final,
            "rise_time": t[np.argmax(moved >= 0.9)]}


def gravity_gap(kp, pose_deg):
    """At rest the motor's push kp × gap balances gravity, so gap = gravity torque / kp (radians)."""
    model = make_model(kp, 0.0)
    data = mujoco.MjData(model)
    data.qpos[:] = np.radians(pose_deg)
    mujoco.mj_forward(model, data)
    torque = data.qfrc_bias[0]                     # with the arm at rest, this is gravity's torque on the shoulder
    return torque, torque / kp


if __name__ == "__main__":
    print("Step response, shoulder 90° → 30° (elbow target stays 0°):")
    for name, (kp, kv) in GAINS.items():
        rows = step_response(kp, kv)
        s = summary(rows, 30.0)
        print(f"  {name:15s} kp {kp:3.1f} kv {kv:4.2f}: first push {rows[0][3]:+.4f} N·m, furthest {s['furthest']:.1f}°, after 2 s {s['final']:.1f}° "
              f"(gap {s['gap']:.1f}°, elbow {rows[-1][2]:.1f}°), 90% of the way after {s['rise_time']:.3f} s")
        if kv > 0:
            torque, gap = gravity_gap(kp, rows[-1][1:3])
            print(f"      by hand: gravity torque {torque:.4f} N·m / kp {kp} = {gap:.3g} rad = {math.degrees(gap):.1f}°")

    print("\nTeleoperation at 30 Hz (strong P + D), shoulder column, degrees:")
    frames, _ = teleoperate(*GAINS["strong P + D"])
    for f in frames[15:66:5]:
        print(f"  frame {f['frame_index']:3d}  t {f['timestamp']:.3f} s  observation.state {f['observation.state'][0]:6.1f}"
              f"  action {f['action'][0]:6.1f}")

    print("\nCalibration: raw counts to degrees (range 1010 … 3070)")
    for raw in (1010, 2040, 2500, 3070):
        print(f"  raw {raw} → {raw_to_degrees(raw, 1010, 3070):.2f}°")