"""The grid arm's task in MuJoCo: the same two 10 cm links, barrier and target, now with motors, gravity, thickness
and a solid wall. Touching the wall ends the episode as a failure, as it would damage a real arm.

Run: python3 classes/ml/labs/robot-arm/mujoco_arm.py

The task is a Gymnasium environment (Gymnasium basics). Its commands are the grid's four: shoulder or elbow ±10°.
"""
import gymnasium as gym
import mujoco
import numpy as np
from gymnasium import spaces

import world as W

PERIOD = 0.2          # seconds of physics per command
CHECK = 5             # physics steps between wall checks (5 ms)

XML = """
<mujoco model="two-link arm and wall">
  <option timestep="0.001" gravity="0 0 -9.81"/>
  <worldbody>
    <geom name="floor" type="plane" size="0.3 0.1 0.01" pos="0 0 -0.006"/>
    <geom name="wall" type="box" pos="0.0975 0 0.025" size="0.0075 0.05 0.025"/>
    <body name="upper_arm">
      <joint name="shoulder" type="hinge" axis="0 -1 0" armature="0.02" range="-5 185"/>
      <geom type="capsule" fromto="0 0 0 0.1 0 0" size="0.004" mass="0.20"/>
      <body name="forearm" pos="0.1 0 0">
        <joint name="elbow" type="hinge" axis="0 -1 0" armature="0.02" range="-185 5"/>
        <geom type="capsule" fromto="0 0 0 0.1 0 0" size="0.004" mass="0.15"/>
        <site name="hand" pos="0.1 0 0"/>
      </body>
    </body>
  </worldbody>
  <contact><exclude body1="upper_arm" body2="forearm"/></contact>
  <actuator>
    <position joint="shoulder" kp="6" kv="0.4"/>
    <position joint="elbow" kp="6" kv="0.4"/>
  </actuator>
</mujoco>
"""


def pose(cell):
    """Grid cell (row, column) → joint angles in degrees."""
    return np.array([W.ANGLES[cell[0]], W.ELBOW[cell[1]]], dtype=float)


def cell(angles):
    """Joint angles → the nearest grid cell: how a grid policy reads a measured pose."""
    return int(np.clip(round(angles[0] / 10), 0, 18)), int(np.clip(round((angles[1] + 180) / 10), 0, 18))


class WallReach(gym.Env):
    """Reach the grid's target (14, 4) cm past the wall. Observation: 2 angles (°), 2 speeds (°/s). Action: 4 commands."""

    def __init__(self):
        self.model = mujoco.MjModel.from_xml_string(XML)
        self.data = mujoco.MjData(self.model)
        self.wall = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "wall")
        self.observation_space = spaces.Box(low=np.array([-5, -185, -360, -360], dtype=np.float32),
                                            high=np.array([185, 5, 360, 360], dtype=np.float32))
        self.action_space = spaces.Discrete(4)
        self.trace = None                                     # set to a list to record the pose every 20 ms

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        start = options["start"] if options else STARTS[self.np_random.integers(len(STARTS))]
        mujoco.mj_resetData(self.model, self.data)
        self.command = pose(start)
        self.data.qpos[:] = self.data.ctrl[:] = np.radians(self.command)
        mujoco.mj_forward(self.model, self.data)
        return self.observe(), {}

    def step(self, action):
        self.command = np.clip(self.command + W.ACTIONS[action] * 10, [0, -180], [180, 0])
        self.data.ctrl[:] = np.radians(self.command)
        touched = False
        for _ in range(round(PERIOD / self.model.opt.timestep) // CHECK):
            mujoco.mj_step(self.model, self.data, nstep=CHECK)
            touched = touched or self.touching_wall()
            if self.trace is not None and round(self.data.time * 1000) % 20 == 0:
                self.trace.append(np.degrees(self.data.qpos).round(2).tolist())
        reached = bool(np.linalg.norm(self.hand() - W.TARGET) < 1.5)
        terminated = touched or reached                       # a touch ends the episode as a failure
        reward = 1.0 if reached and not touched else 0.0
        return self.observe(), reward, terminated, False, {"touched": touched, "reached": reached and not touched}

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

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

    def touching_wall(self):
        return any(self.wall in (c.geom1, c.geom2) for c in self.data.contact[:self.data.ncon])


gym.register("WallReach-v0", entry_point=WallReach, max_episode_steps=60)


def touches_at_rest(start):
    env = WallReach()
    env.reset(options={"start": start})
    return env.touching_wall()


STARTS = [s for s in W.STARTS if not touches_at_rest(s)]   # the grid's starts that a 4 mm-thick arm can stand in


def run(policy, start, env=None):
    """One episode from `start`. Returns success, commands used, whether the wall was touched, and the poses."""
    env = env or gym.make("WallReach-v0")
    observation, _ = env.reset(options={"start": start})
    poses = [observation[:2].tolist()]
    for n in range(1, 61):
        observation, reward, terminated, truncated, info = env.step(policy(observation))
        poses.append(observation[:2].tolist())
        if terminated or truncated:
            return {"success": info["reached"], "commands": n, "touched": info["touched"], "poses": poses}


def table_policy(q):
    """Act from a 19 × 19 × 4 table: read the measured pose as its nearest grid cell."""
    return lambda observation: int(np.argmax(q[cell(observation[:2])]))


def careful_planner(margin=1.0):
    """The grid's full-model planner, with the barrier made `margin` cm bigger on every side: a person keeping clear."""
    blocked = np.array([[forbidden(a, b, margin) for b in W.ELBOW] for a in W.ANGLES])
    goal = W.GOAL & ~blocked
    values = np.zeros((19, 19))
    for _ in range(200):
        q = np.zeros((19, 19, 4))
        for i in range(19):
            for j in range(19):
                for a, (di, dj) in enumerate(W.ACTIONS):
                    after = (i + di, j + dj)
                    if not (0 <= after[0] < 19 and 0 <= after[1] < 19) or blocked[after]:
                        after = (i, j)
                    q[i, j, a] = 1.0 if goal[after] else W.GAMMA * values[after]
        values = q.max(axis=2)
    return q


def forbidden(shoulder, elbow, margin):
    joint, hand = W.forward(shoulder, elbow)
    fraction = np.linspace(0, 1, 60)[:, None]
    points = np.concatenate([fraction * joint, joint + fraction * (hand - joint)])
    near_wall = ((points[:, 0] >= 9 - margin) & (points[:, 0] <= 10.5 + margin)
                 & (points[:, 1] >= 0) & (points[:, 1] <= 5 + margin)).any()
    return bool((points[:, 1] < -1e-9).any() or near_wall)


def q_learning(seed, episodes=300):
    """The Q-table workshop's learner, now trying inside MuJoCo from home. A wall touch ends the try with nothing."""
    rng = np.random.default_rng(seed)
    q = np.zeros((19, 19, 4))
    env = gym.make("WallReach-v0", max_episode_steps=400)
    commands = crashes = 0
    first = None
    for episode in range(episodes):
        observation, _ = env.reset(options={"start": W.HOME})
        state = cell(observation[:2])
        while True:
            action = int(rng.integers(4)) if rng.random() < 0.2 else int(np.argmax(q[state] + 1e-9 * rng.random(4)))
            observation, reward, terminated, truncated, info = env.step(action)
            after = cell(observation[:2])
            target = reward + (0 if terminated else W.GAMMA * q[after].max())
            q[state][action] += 0.5 * (target - q[state][action])
            commands += 1
            crashes += info["touched"]
            state = after
            if info["reached"] and first is None:
                first = episode + 1
            if terminated or truncated:
                break
    return {"q": q, "commands": commands, "crashes": crashes, "first_success": first}


def demonstrate(starts, expert):
    """Record the expert driving the arm in MuJoCo: (measured pose as a grid position, command) at every step."""
    pairs, commands = [], 0
    env = gym.make("WallReach-v0")
    for start in starts:
        observation, _ = env.reset(options={"start": start})
        while True:
            action = int(np.argmax(expert[cell(observation[:2])]))
            pairs.append(((observation[0] / 10, (observation[1] + 180) / 10), action))
            observation, reward, terminated, truncated, info = env.step(action)
            commands += 1
            if terminated or truncated:
                break
    return pairs, commands


def network_policy(policy):
    """Act from a trained 2 → 16 → 4 network, fed the measured angles (not rounded to the grid)."""
    import torch
    def act(observation):
        with torch.no_grad():
            x = torch.tensor(W.encode([[observation[0] / 10, (observation[1] + 180) / 10]]))
            return int(policy(x).argmax())
    return act


if __name__ == "__main__":
    print(f"{len(W.STARTS)} grid starts; {len(W.STARTS) - len(STARTS)} touch the wall at rest; {len(STARTS)} remain")
    for name, q in (("grid planner", W.reference()), ("careful planner", careful_planner())):
        runs = [run(table_policy(q), s) for s in STARTS]
        print(f"{name:16s} {sum(r['success'] for r in runs)}/{len(STARTS)} reach the target; "
              f"{sum(r['touched'] for r in runs)} touch the wall")

    import imitation
    learned = q_learning(seed=0)
    runs = [run(table_policy(learned["q"]), s) for s in STARTS]
    print(f"learning by trying (seed 0, 300 tries from home): {learned['commands']} commands, {learned['crashes']} wall crashes; "
          f"home {'reached' if run(table_policy(learned['q']), W.HOME)['success'] else 'failed'}, {sum(r['success'] for r in runs)}/{len(STARTS)} starts")

    rng = np.random.default_rng(0)
    pairs, commands = demonstrate([STARTS[i] for i in rng.choice(len(STARTS), 10, replace=False)], careful_planner())
    policy = imitation.train(pairs, seed=0)
    runs = [run(network_policy(policy), s) for s in STARTS]
    print(f"learning by watching (10 demonstrations, {commands} commands): {sum(r['success'] for r in runs)}/{len(STARTS)} starts, "
          f"{sum(r['touched'] for r in runs)} touch the wall")