"""
Arm IK + scripted pick-and-place teacher for the demo2 arm+Shadow-Hand.

The teacher is the privileged demonstrator for behaviour cloning — NOT the
deployed controller (the runtime is a learned model, per the project's
model-driven-control preference). It scripts a clean pick→place using:

  - ArmIK: damped-least-squares Jacobian IK on the 5 arm joints toward the palm
    grasp site (algorithm ported from models/rx1_ik.py).
  - PickPlaceTeacher: phase machine approach→lower→close→weld→lift→carry→place
    →release (sequence ported from models/expert.py), with a hand↔object weld
    captured on grasp so the object is carried.
"""
from __future__ import annotations

from enum import Enum, auto
import numpy as np
import mujoco

# arm actuators / joints (first 5), grasp site, weld
ARM_ACT = [0, 1, 2, 3, 4]
ARM_JOINTS = ["waist_yaw", "shoulder_pitch", "shoulder_roll", "elbow", "wrist_pitch"]

# finger actuator targets (names match the merged Shadow-Hand actuators)
_FINGERS = ("FF", "MF", "RF", "LF")
_OPEN = {f"rh_A_{f}J3": 0.0 for f in _FINGERS}
_OPEN.update({f"rh_A_{f}J0": 0.0 for f in _FINGERS})
_OPEN.update({"rh_A_THJ4": 0.6, "rh_A_THJ2": 0.0, "rh_A_THJ1": 0.0})
_CLOSED = {f"rh_A_{f}J3": 1.35 for f in _FINGERS}
_CLOSED.update({f"rh_A_{f}J0": 2.4 for f in _FINGERS})
_CLOSED.update({"rh_A_THJ4": 1.1, "rh_A_THJ2": 0.5, "rh_A_THJ1": 1.0})


def _name_ids(model, names, objtype):
    return [mujoco.mj_name2id(model, objtype, n) for n in names]


class ArmIK:
    """Damped-least-squares IK for the arm's palm grasp site."""

    def __init__(self, model):
        self.model = model
        self.nu = model.nu
        self._lo = model.actuator_ctrlrange[:, 0].copy()
        self._hi = model.actuator_ctrlrange[:, 1].copy()
        self._site = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "grasp")
        jids = _name_ids(model, ARM_JOINTS, mujoco.mjtObj.mjOBJ_JOINT)
        self._dofs = [int(model.jnt_dofadr[j]) for j in jids]
        self._qpos = [int(model.jnt_qposadr[j]) for j in jids]
        self._fopen = self._finger_vec(_OPEN)
        self._fclosed = self._finger_vec(_CLOSED)
        self._aid = {mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i): i
                     for i in range(model.nu)}
        self._palm = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "rh_palm")
        # The "approach axis" in the palm's local frame (palm centre → fingertips).
        # We constrain this to point straight down (world -z) so the hand always
        # picks top-down instead of coming in sideways (position-only IK left the
        # hand orientation uncontrolled → it grasped sideways / looked upside down).
        self._approach_local = self._measure_approach_axis()

    def _measure_approach_axis(self):
        d = mujoco.MjData(self.model)
        mujoco.mj_kinematics(self.model, d); mujoco.mj_comPos(self.model, d)
        R = d.xmat[self._palm].reshape(3, 3)
        palm = d.xpos[self._palm]
        ff = d.xpos[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "rh_ffdistal")]
        a_world = ff - palm
        a_local = R.T @ a_world
        return a_local / (np.linalg.norm(a_local) + 1e-9)

    def _finger_vec(self, spec):
        v = {}
        for k, val in spec.items():
            i = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, k)
            if i >= 0:
                v[i] = val
        return v

    def rest(self):
        """Rest ctrl: arm reaching forward-down over the table (a good IK basin),
        fingers open."""
        q = np.zeros(self.nu, np.float32)
        q[1] = -0.9      # shoulder_pitch (reach forward-down)
        q[3] = -0.6      # elbow (mostly extended)
        q[4] = -0.5      # wrist_pitch (point fingers down)
        for i, val in self._fopen.items():
            q[i] = val
        return np.clip(q, self._lo, self._hi)

    def set_fingers(self, q, closed):
        q = q.copy()
        for i, val in (self._fclosed if closed else self._fopen).items():
            q[i] = val
        return np.clip(q, self._lo, self._hi)

    def solve(self, target_xyz, q_ctrl, n_iter=300, lr=0.5, lam=0.06,
              down=True, w_rot=0.4):
        """Return a full ctrl vector with only the arm actuators updated so the
        grasp site reaches target_xyz. With down=True, also orient the hand so
        its approach axis points straight down (top-down pick) — a combined
        position + orientation objective on the 5-DOF arm (3 DOF reach the
        point, the remaining 2 aim the hand down)."""
        d = mujoco.MjData(self.model)
        d.qpos[:] = 0.0
        q_arm = np.array([q_ctrl[a] for a in ARM_ACT], float)
        lo = np.array([self._lo[a] for a in ARM_ACT])
        hi = np.array([self._hi[a] for a in ARM_ACT])
        target = np.asarray(target_xyz, float)
        down_vec = np.array([0.0, 0.0, -1.0])
        for _ in range(n_iter):
            for k, qi in enumerate(self._qpos):
                d.qpos[qi] = q_arm[k]
            mujoco.mj_kinematics(self.model, d)
            mujoco.mj_comPos(self.model, d)
            err_p = target - d.site_xpos[self._site]
            jacp = np.zeros((3, self.model.nv))
            jacr = np.zeros((3, self.model.nv)) if down else None
            mujoco.mj_jacSite(self.model, d, jacp, jacr, self._site)
            if down:
                R = d.xmat[self._palm].reshape(3, 3)
                a_world = R @ self._approach_local
                # angular error that rotates the approach axis toward straight-down
                err_r = np.cross(a_world, down_vec)
                err = np.concatenate([err_p, w_rot * err_r])
                J = np.vstack([jacp[:, self._dofs], w_rot * jacr[:, self._dofs]])
                damp = lam ** 2 * np.eye(6)
            else:
                if np.linalg.norm(err_p) < 0.004:
                    break
                err = err_p
                J = jacp[:, self._dofs]
                damp = lam ** 2 * np.eye(3)
            if down and np.linalg.norm(err_p) < 0.004 and np.linalg.norm(err_r) < 0.03:
                break
            dq = J.T @ np.linalg.solve(J @ J.T + damp, err)
            q_arm = np.clip(q_arm + lr * dq, lo, hi)
        out = np.asarray(q_ctrl, np.float32).copy()
        for k, a in enumerate(ARM_ACT):
            out[a] = q_arm[k]
        return out


class _Phase(Enum):
    APPROACH = auto(); LOWER = auto(); CLOSE = auto(); WELD = auto()
    LIFT = auto(); CARRY = auto(); PLACE = auto(); SETTLE = auto()
    RELEASE = auto(); DONE = auto()


# dwell phases run a fixed number of steps; reach phases run until the grasp
# site converges on the waypoint AND a minimum dwell passes (so the arm settles
# and the carried object stops swinging before the next waypoint), capped so a
# slewed arm under load can't stall forever.
_DWELL = {_Phase.CLOSE: 60, _Phase.WELD: 3, _Phase.SETTLE: 80, _Phase.RELEASE: 40}
_REACH_TOL = 0.03
_REACH_MIN = 50          # min steps in a reach phase before advancing
_REACH_CAP = 220
_PRE_Z, _LIFT_Z, _PLACE_DZ = 0.10, 0.10, 0.005


class PickPlaceTeacher:
    """Scripted pick→place demonstrator (BC teacher). Reach phases converge on
    their waypoint before advancing; dwell phases run a fixed number of steps."""

    def __init__(self, ik: ArmIK, pick_xyz, place_xyz):
        self.ik = ik
        self.nu = ik.nu
        self._pick = np.asarray(pick_xyz, float)
        self._place = np.asarray(place_xyz, float)
        self._ctrl = ik.rest()
        self._phase = _Phase.APPROACH
        self._k = 0
        self._goal = None          # cached ctrl solution for the current reach phase

    @property
    def done(self):
        return self._phase == _Phase.DONE

    @property
    def phase(self):
        return self._phase.name.lower()

    def _advance(self):
        self._k = 0
        self._goal = None
        order = list(_Phase)
        self._phase = order[min(order.index(self._phase) + 1, len(order) - 1)]

    def act(self, q_cur, obj_live, site_xyz):
        """Return (ctrl[nu], grasp_bool).
        obj_live = object xyz; site_xyz = current grasp-site xyz (for convergence).
        """
        obj = np.asarray(obj_live, float)
        grasp = self._phase in (_Phase.WELD, _Phase.LIFT, _Phase.CARRY,
                                _Phase.PLACE, _Phase.SETTLE)
        p = self._phase
        self._k += 1

        # reach phases: set a waypoint, drive toward it, advance when reached
        reach_target = None
        if p == _Phase.APPROACH:
            reach_target = self._pick + [0, 0, _PRE_Z]
        elif p == _Phase.LOWER:
            reach_target = obj.copy()
        elif p == _Phase.LIFT:
            reach_target = self._pick + [0, 0, _LIFT_Z]
        elif p == _Phase.CARRY:
            reach_target = self._place + [0, 0, _LIFT_Z]
        elif p == _Phase.PLACE:
            reach_target = self._place + [0, 0, _PLACE_DZ]

        if reach_target is not None:
            # solve ONCE per phase and hold the goal — re-solving from the
            # mid-slew pose every step makes the target oscillate and never
            # converge.
            if self._goal is None:
                self._goal = self.ik.solve(reach_target, q_cur)
            self._ctrl = self.ik.set_fingers(self._goal, grasp)
            reached = np.linalg.norm(np.asarray(site_xyz) - reach_target) < _REACH_TOL
            if (reached and self._k >= _REACH_MIN) or self._k >= _REACH_CAP:
                self._advance()
        else:
            # dwell phases: close / weld / settle / release, fixed steps
            if p in (_Phase.CLOSE, _Phase.WELD, _Phase.SETTLE):
                self._ctrl = self.ik.set_fingers(q_cur, grasp)
            elif p == _Phase.RELEASE:
                self._ctrl = self.ik.set_fingers(q_cur, False)
            if self._k >= _DWELL.get(p, 1):
                self._advance()
        return self._ctrl, grasp