"""
Shadow Dexterous Hand controller — /demo2 (5-finger gripper).

Maps a natural-language gesture name to position targets for the 20 Shadow Hand
actuators (thumb TH + first/middle/ring/little fingers + 2 wrist joints). Static
gestures are fixed poses; `wave` and `count` are time-varying finger gaits.

Joint convention (from the Menagerie model):
  *J4  finger spread/abduction (±0.35)      *J3  proximal flexion (0 open .. 1.57)
  *J0  coupled mid+distal flexion (0 .. 3.14)
  TH:  THJ4 opposition (0..1.22), THJ2/THJ1 flexion
This is a self-contained demo stack; nothing here depends on the RX1 v1/v2 code.
"""
from __future__ import annotations

import math

import numpy as np
import mujoco

# flexion presets
_F3, _F0 = 1.45, 2.9          # proximal / coupled-distal "closed"
_TH = {"rh_A_THJ4": 1.1, "rh_A_THJ2": 0.55, "rh_A_THJ1": 1.1}   # thumb across palm

def _flex(*fingers):
    d = {}
    for f in fingers:
        d[f"rh_A_{f}J3"] = _F3
        d[f"rh_A_{f}J0"] = _F0
    return d

def _ext(*fingers):
    d = {}
    for f in fingers:
        d[f"rh_A_{f}J3"] = 0.0
        d[f"rh_A_{f}J0"] = 0.0
    return d

# ── static gesture pose library (actuator-name → target; unset → 0) ──────────
GESTURES = {
    "open":   {},                                              # flat hand
    "rest":   {**_flex("FF", "MF", "RF", "LF"), **_TH},        # relaxed light curl→ overwritten below
    "fist":   {**_flex("FF", "MF", "RF", "LF"), **_TH},        # full grasp
    "grasp":  {**_flex("FF", "MF", "RF", "LF"), **_TH},        # alias of fist
    "point":  {**_ext("FF"), **_flex("MF", "RF", "LF"), **_TH},        # index out
    "peace":  {**_ext("FF", "MF"), **_flex("RF", "LF"), **_TH},        # V sign
    "three":  {**_ext("FF", "MF", "RF"), **_flex("LF"), **_TH},
    "pinch":  {"rh_A_THJ4": 1.1, "rh_A_THJ2": 0.4, "rh_A_THJ1": 0.7,
               "rh_A_FFJ3": 0.75, "rh_A_FFJ0": 0.6,
               **_flex("MF", "RF", "LF")},                     # thumb-index pinch
    "ok":     {"rh_A_THJ4": 1.1, "rh_A_THJ2": 0.35, "rh_A_THJ1": 0.7,
               "rh_A_FFJ3": 0.8, "rh_A_FFJ0": 0.9},            # OK ring, others up
    "spread": {"rh_A_FFJ4": 0.35, "rh_A_MFJ4": 0.1,
               "rh_A_RFJ4": -0.1, "rh_A_LFJ4": -0.35, "rh_A_LFJ5": 0.6},  # splay
}
# rest = a soft curl, lighter than fist
GESTURES["rest"] = {f"rh_A_{f}J3": 0.4 for f in ("FF", "MF", "RF", "LF")}

ALIASES = {"close": "fist", "hold": "grasp", "release": "open", "flat": "open",
           "victory": "peace", "index": "point", "splay": "spread"}

DYNAMIC = ("wave", "count")
_ORDER = ("FF", "MF", "RF", "LF")     # finger order for the counting gait


class DexHandController:
    def __init__(self, model: mujoco.MjModel):
        self.model = model
        self.nu = model.nu
        self._idx = {mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i): i
                     for i in range(model.nu)}
        self._lo = model.actuator_ctrlrange[:, 0].copy()
        self._hi = model.actuator_ctrlrange[:, 1].copy()
        self.gesture = "open"

    def set_gesture(self, name: str):
        n = (name or "").strip().lower()
        for w in n.split():                      # pick the first known keyword
            w = ALIASES.get(w, w)
            if w in GESTURES or w in DYNAMIC:
                self.gesture = w
                return self.gesture
        n = ALIASES.get(n, n)
        self.gesture = n if (n in GESTURES or n in DYNAMIC) else "open"
        return self.gesture

    def _pose(self, spec: dict) -> np.ndarray:
        q = np.zeros(self.nu, dtype=np.float32)
        for name, val in spec.items():
            if name in self._idx:
                q[self._idx[name]] = val
        return np.clip(q, self._lo, self._hi)

    def target(self, t: float) -> np.ndarray:
        """Position targets for the current gesture at time t (seconds)."""
        g = self.gesture
        if g == "wave":
            q = self._pose({})
            for k, f in enumerate(_ORDER):
                phase = 2 * math.pi * 1.0 * t + k * 0.6
                q[self._idx[f"rh_A_{f}J3"]] = 0.7 + 0.7 * (0.5 + 0.5 * math.sin(phase))
            return np.clip(q, self._lo, self._hi)
        if g == "count":
            n = int(t // 0.8) % (len(_ORDER) + 1)          # raise 0..4 fingers
            up = set(_ORDER[:n])
            spec = {}
            for f in _ORDER:
                spec[f"rh_A_{f}J3"] = 0.0 if f in up else _F3
                spec[f"rh_A_{f}J0"] = 0.0 if f in up else _F0
            spec.update(_TH)
            return self._pose(spec)
        return self._pose(GESTURES.get(g, {}))

    @staticmethod
    def known() -> list:
        return sorted(set(GESTURES) | set(DYNAMIC) | set(ALIASES))