"""Gymnasium basics: the Robot basics arm in MuJoCo, wrapped as a Gymnasium task.

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

Gymnasium is the standard way to package a task for a learner: reset() starts an episode, step(action) moves it on.
LeRobot's simulators (gym-aloha, gym-hil) use exactly these calls.
"""

import time

import gymnasium as gym
import mujoco
import mujoco.viewer
import numpy as np
from gymnasium import spaces

L1, L2 = 0.10, 0.10                                  # link lengths in metres, as in Robot basics
COMMANDS = np.array([[10, 0], [-10, 0], [0, 10], [0, -10]])   # degrees added to the shoulder or elbow target
COMMAND_NAMES = ["shoulder +10°", "shoulder −10°", "elbow +10°", "elbow −10°"]
PERIOD = 0.2                                         # seconds of physics per command
START = np.array([90.0, -90.0])                      # shoulder up, forearm pointing forward

XML = f"""
<mujoco model="two-link arm">
  <option timestep="0.001" gravity="0 0 -9.81"/>
  <worldbody>
    <light pos="0 -0.3 0.4" dir="0 0.6 -1"/>
    <camera name="side" pos="0 -0.6 0.03" xyaxes="1 0 0 0 0 1"/>
    <site name="target" type="sphere" size="0.015" rgba="0.95 0.45 0.1 0.6"/>
    <geom type="plane" size="0.4 0.3 0.01" pos="0 0 -0.21" rgba="0.82 0.83 0.8 1" contype="0" conaffinity="0"/>
    <body name="upper_arm">
      <joint name="shoulder" type="hinge" axis="0 -1 0" armature="0.02" range="0 180"/>
      <geom type="capsule" fromto="0 0 0 {L1} 0 0" size="0.004" mass="0.20"/>
      <body name="forearm" pos="{L1} 0 0">
        <joint name="elbow" type="hinge" axis="0 -1 0" armature="0.02" range="-180 0"/>
        <geom type="capsule" fromto="0 0 0 {L2} 0 0" size="0.004" mass="0.15"/>
        <site name="hand" pos="{L2} 0 0"/>
      </body>
    </body>
  </worldbody>
  <actuator>
    <position joint="shoulder" kp="6" kv="0.4"/>
    <position joint="elbow" kp="6" kv="0.4"/>
  </actuator>
</mujoco>
"""


def hand_position(shoulder, elbow):
    """Where the hand is for two joint angles in degrees, in centimetres (x forward, y up)."""
    a, b = np.radians(shoulder), np.radians(shoulder + elbow)
    return 100 * np.array([L1 * np.cos(a) + L2 * np.cos(b), L1 * np.sin(a) + L2 * np.sin(b)])


class ArmReach(gym.Env):
    """Bring the hand within 1.5 cm of a target. One step is one command, then 0.2 s of physics."""
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 5}

    def __init__(self, render_mode=None):
        self.render_mode = render_mode               # "human": MuJoCo's viewer window; "rgb_array": pictures
        self.viewer = self.renderer = None
        self.model = mujoco.MjModel.from_xml_string(XML)
        self.data = mujoco.MjData(self.model)
        # What the learner sees: 2 joint angles (°), 2 joint speeds (°/s), the target's x and y (cm).
        self.observation_space = spaces.Box(low=np.array([-5, -185, -360, -360, -20, -20], dtype=np.float32),
                                            high=np.array([185, 5, 360, 360, 20, 20], dtype=np.float32))
        # What the learner may do: one of the 4 commands.
        self.action_space = spaces.Discrete(4)

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)                     # seeds self.np_random: the same seed, the same target
        shoulder = self.np_random.integers(1, 18) * 10            # a target the arm can reach: the hand of a random pose
        elbow = -self.np_random.integers(1, 18) * 10
        self.target = hand_position(shoulder, elbow)
        self.model.site("target").pos = [self.target[0] / 100, 0, self.target[1] / 100]   # the orange ball, for pictures
        mujoco.mj_resetData(self.model, self.data)
        self.command = START.copy()
        self.data.qpos[:] = self.data.ctrl[:] = np.radians(self.command)
        mujoco.mj_forward(self.model, self.data)
        if self.render_mode == "human":
            self.render()
        return self.observe(), {}

    def step(self, action):
        self.command = np.clip(self.command + COMMANDS[action], [0, -180], [180, 0])
        self.data.ctrl[:] = np.radians(self.command)                 # the new target angles for the motors
        mujoco.mj_step(self.model, self.data, nstep=round(PERIOD / self.model.opt.timestep))
        distance = np.linalg.norm(self.hand() - self.target)
        terminated = bool(distance < 1.5)            # the task is done
        reward = 1.0 if terminated else 0.0
        if self.render_mode == "human":
            self.render()
        return self.observe(), reward, terminated, False, {"distance": distance}

    def render(self):
        """The picture comes from MuJoCo: Gymnasium only defines when and how to ask for it."""
        if self.render_mode == "rgb_array":
            if self.renderer is None:
                self.renderer = mujoco.Renderer(self.model, 240, 320)
            self.renderer.update_scene(self.data, camera="side")
            return self.renderer.render()
        if self.render_mode == "human":
            if self.viewer is None:
                self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
                self.viewer.cam.lookat[:], self.viewer.cam.distance = [0, 0, 0.02], 0.65     # look at the arm
                self.viewer.cam.azimuth, self.viewer.cam.elevation = 90, -8
            self.viewer.sync()
            time.sleep(1 / self.metadata["render_fps"])

    def close(self):
        if self.viewer is not None:
            self.viewer.close()
            time.sleep(0.5)                          # let the viewer's window thread finish before Python exits
        if self.renderer is not None:
            self.renderer.close()

    def observe(self):
        angles, speeds = np.degrees(self.data.qpos), np.degrees(self.data.qvel)
        return np.concatenate([angles, speeds, self.target]).astype(np.float32)

    def hand(self):
        x, _, z = self.data.site("hand").xpos
        return 100 * np.array([x, z])


gym.register("ArmReach-v0", entry_point=ArmReach, max_episode_steps=40)   # gym.make adds a 40-command time limit


def closest_policy(observation):
    """A hand-written policy: the command whose pose puts the hand nearest the target."""
    angles, target = observation[:2], observation[4:]
    return int(np.argmin([np.linalg.norm(hand_position(*(angles + c)) - target) for c in COMMANDS]))


def random_policy(env):
    return lambda observation: int(env.action_space.sample())


def run_episode(env, policy, seed):
    """The loop every Gymnasium learner uses."""
    observation, info = env.reset(seed=seed)
    steps = []
    while True:
        action = policy(observation)
        observation, reward, terminated, truncated, info = env.step(action)
        steps.append((action, observation, reward, terminated, truncated))
        if terminated or truncated:
            return steps


def success_rate(policy_for, seeds):
    env = gym.make("ArmReach-v0")
    env.action_space.seed(0)                         # the random policy's draws repeat too
    results = [run_episode(env, policy_for(env), seed) for seed in seeds]
    return sum(r[-1][3] for r in results), [len(r) for r in results]


if __name__ == "__main__":
    env = gym.make("ArmReach-v0")
    print("observation space:", env.observation_space, "| action space:", env.action_space)

    observation, info = env.reset(seed=3)
    print("reset(seed=3):", observation.round(2))
    print("reset(seed=3) again:", env.reset(seed=3)[0].round(2), "| reset(seed=4):", env.reset(seed=4)[0].round(2))

    env.reset(seed=3)
    observation, reward, terminated, truncated, info = env.step(0)
    print("step(0):", observation.round(2), reward, terminated, truncated, f"distance {info['distance']:.2f} cm")

    from gymnasium.utils.env_checker import check_env
    check_env(gym.make("ArmReach-v0").unwrapped)
    print("check_env: passed")

    for name, policy_for in (("random", random_policy), ("closest", lambda env: closest_policy)):
        reached, lengths = success_rate(policy_for, range(100))
        print(f"{name:8s} policy: {reached}/100 targets reached; commands per episode: median {np.median(lengths):.0f}")

    envs = gym.make_vec("ArmReach-v0", num_envs=10, vectorization_mode="sync")
    observations, infos = envs.reset(seed=1000)                  # copy i gets seed 1000 + i
    print("10 copies: observations", observations.shape, "| first target of copy 1:", observations[1, 4:].round(2),
          "| alone with seed 1001:", env.reset(seed=1001)[0][4:].round(2))
    observations, rewards, terminated, truncated, infos = envs.step(np.zeros(10, dtype=int))
    print("one step of all 10:", observations.shape, rewards.shape, terminated.shape)