"""
Privileged-state pick-and-place EXPERT  (v2 — demonstration teacher only).

This is NOT the runtime controller. It generates clean demonstrations the
learned arms-only action model (models/act_model.py, trained by
scripts/train_act.py) imitates; at deployment the *learned* model drives the
robot and this teacher is never loaded.

Control regime (must match models/policy_learned.py exactly): the non-arm
joints (torso, head, the 18 fingers, the unused left arm) are commanded to
their CURRENT pose — i.e. held compliant — while the right-arm IK and a grasp
bit do the work. This matters: the RX1's arm alone cannot reach the table from
a bolt-upright torso (the shoulder cantilever droops ~0.4 rad); leaving the
torso compliant lets it lean forward under load and the hand reaches. Both the
teacher and the runtime use this same compliant regime so the cloned arm
targets behave identically at train and test (verified: replaying recorded
actions reproduces placement to <0.06 m).

Deterministic and STEP-COUNT-driven so offline collection runs fast. Per step:
    joints(arm 14) — right(7)+left(7) arm position targets (left ≈ held)
    grasp(1)       — 1 while the cube should be welded to the hand
"""
from __future__ import annotations

from enum import Enum, auto

import numpy as np
import mujoco

from models.rx1_ik import RX1IKSolver


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


_STEPS = {
    _Phase.APPROACH: 40, _Phase.LOWER: 35, _Phase.CLOSE: 15, _Phase.WELD: 3,
    _Phase.LIFT: 35, _Phase.CARRY: 35, _Phase.PLACE: 35, _Phase.RELEASE: 15,
}
_PRE_Z, _LIFT_Z, _PLACE_DZ = 0.12, 0.20, 0.006


class PickPlaceExpert:
    """Step-driven privileged pick→place teacher (compliant non-arm regime)."""

    ACTION_DIM = 41   # 40 joint targets + 1 grasp bit (collector keeps arm 14 + grasp)

    def __init__(self, ik: RX1IKSolver, pick_pos, place_pos):
        self._ik    = ik
        self.nu     = ik.nu
        self._pick  = np.asarray(pick_pos, float).copy()
        self._place = np.asarray(place_pos, float).copy()
        self._offset = np.zeros(3)
        self._ctrl  = ik.rest()
        self._phase = _Phase.APPROACH
        self._k = 0

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

    def _site(self, q):
        d = mujoco.MjData(self._ik.model)
        d.qpos[: self.nu] = q
        mujoco.mj_kinematics(self._ik.model, d)
        return d.site_xpos[self._ik._rsite].copy()

    def _reach(self, cube_centre_target, q_cur, closed):
        # solve() copies non-arm joints from q_cur ⇒ they stay at the current
        # (compliant) pose; only the right arm is driven toward the target.
        site_target = np.asarray(cube_centre_target, float) - self._offset
        q = self._ik.solve(site_target, q_cur, side="right")
        return self._ik.set_fingers(q, closed=closed, side="right")

    def act(self, q_cur, cube_live):
        cube_live = np.asarray(cube_live, float)
        grasp = self._phase in (_Phase.WELD, _Phase.LIFT, _Phase.CARRY, _Phase.PLACE)

        if self._phase == _Phase.APPROACH:
            self._ctrl = self._reach(self._pick + [0, 0, _PRE_Z], q_cur, False)
        elif self._phase == _Phase.LOWER:
            self._ctrl = self._reach(cube_live, q_cur, False)
        elif self._phase == _Phase.CLOSE:
            self._ctrl = self._ik.set_fingers(q_cur.copy(), True, "right")
        elif self._phase == _Phase.WELD:
            self._offset = cube_live - self._site(q_cur)
        elif self._phase == _Phase.LIFT:
            self._ctrl = self._reach(self._pick + [0, 0, _LIFT_Z], q_cur, True)
        elif self._phase == _Phase.CARRY:
            self._ctrl = self._reach(self._place + [0, 0, _LIFT_Z], q_cur, True)
        elif self._phase == _Phase.PLACE:
            self._ctrl = self._reach(self._place + [0, 0, _PLACE_DZ], q_cur, True)
        elif self._phase == _Phase.RELEASE:
            self._ctrl = self._ik.set_fingers(q_cur.copy(), False, "right")

        action = np.zeros(self.ACTION_DIM, dtype=np.float32)
        action[: self.nu] = self._ctrl
        action[self.nu] = 1.0 if grasp else 0.0
        self._advance()
        return action, grasp

    def _advance(self):
        if self._phase == _Phase.DONE:
            return
        self._k += 1
        if self._k >= _STEPS.get(self._phase, 0):
            self._k = 0
            order = list(_Phase)
            self._phase = order[order.index(self._phase) + 1]