"""
Grasp planning for the /demo2 Shadow Dexterous Hand.

Pipeline: perceive → plan → (executed by the server as preshape → close-until-
contact → hold). This is the planning brain the gesture demo lacked: it looks at
the actual object, decides an approach and an opposition grasp sized to the
object, and reports the decision.

Method — pragmatic opposition grasp inspired by Liu, Jiang & Cheng, "A Fast Grasp
Planning Algorithm for Humanoid Robot Hands" (Biomimetics 2024, 9(10), 599). The
paper's idea is to reach a multi-finger force-closure grasp cheaply instead of
solving the full nonlinear/convex-hull optimisation. We approximate that with a
thumb-opposing-fingers wrap, sized to the object's graspable width, then close
each finger until it contacts the surface — opposition + multiple wrap contacts
give an (approximate) force-closure hold without the expensive search.

Ground-truth perception: object pose/size are read straight from MuJoCo (the sim
is treated as a perfect sensor), so the focus stays on the planning loop.
"""
from __future__ import annotations

import numpy as np
import mujoco

FINGERS = ("FF", "MF", "RF", "LF")
TIP_BODY = {"FF": "rh_ffdistal", "MF": "rh_mfdistal", "RF": "rh_rfdistal",
            "LF": "rh_lfdistal", "TH": "rh_thdistal"}

# closed-grasp joint extents (rad), scaled by the per-object strength
_J3_CLOSED = 1.45     # proximal flexion
_J0_CLOSED = 2.60     # coupled mid+distal flexion


def _bid(model, name):
    return mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)


# a finger "contacts" the object if ANY of its links (knuckle→distal) touch it,
# not just the tip — a power grasp often seats the object against the proximal
# links / thumb base, so tip-only detection misses real contacts.
_FINGER_PREFIX = {"FF": "rh_ff", "MF": "rh_mf", "RF": "rh_rf",
                  "LF": "rh_lf", "TH": "rh_th"}


class GraspMaps:
    """Precomputed geom↔finger lookups for fast object-contact detection."""

    def __init__(self, model, obj_body="object"):
        self.obj_geoms = set(self._geoms_of(model, _bid(model, obj_body)))
        self.finger_geoms = {f: set() for f in _FINGER_PREFIX}
        for gi in range(model.ngeom):
            bn = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY,
                                   model.geom_bodyid[gi]) or ""
            for f, pre in _FINGER_PREFIX.items():
                if bn.startswith(pre):
                    self.finger_geoms[f].add(gi)

    @staticmethod
    def _geoms_of(model, bid):
        return [gi for gi in range(model.ngeom) if model.geom_bodyid[gi] == bid]

    def fingertip_contacts(self, model, data):
        """Set of finger keys currently touching the object (any link)."""
        touch = set()
        for ci in range(data.ncon):
            g1, g2 = data.contact[ci].geom1, data.contact[ci].geom2
            in_obj = (g1 in self.obj_geoms, g2 in self.obj_geoms)
            if not any(in_obj):
                continue
            other = g2 if in_obj[0] else g1
            for f, fg in self.finger_geoms.items():
                if other in fg:
                    touch.add(f)
        return touch


def set_weld(model, data, eq_name="grasp", active=True, body2=None):
    """Toggle the hand↔object weld, capturing the current relative pose on the
    rising edge so the object holds exactly where the fingers seated it.

    body2: optionally repoint the weld to a different object body before
    activating (so one weld can hold whichever object is being picked)."""
    eq_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, eq_name)
    if eq_id < 0:
        return
    if body2 is not None:
        bid = _bid(model, body2) if isinstance(body2, str) else int(body2)
        if bid >= 0:
            model.eq_obj2id[eq_id] = bid
    was = bool(data.eq_active[eq_id])
    if active and not was:
        b1, b2 = model.eq_obj1id[eq_id], model.eq_obj2id[eq_id]
        p1, r1 = data.xpos[b1], data.xmat[b1].reshape(3, 3)
        p2, r2 = data.xpos[b2], data.xmat[b2].reshape(3, 3)
        relpos = r1.T @ (p2 - p1)
        quat = np.zeros(4)
        mujoco.mju_mat2Quat(quat, (r1.T @ r2).reshape(9))
        model.eq_data[eq_id, 0:3] = 0.0
        model.eq_data[eq_id, 3:6] = relpos
        model.eq_data[eq_id, 6:10] = quat
    data.eq_active[eq_id] = 1 if active else 0


def perceive(model, data, obj_body="object"):
    """Ground-truth object + hand geometry for planning."""
    oid = _bid(model, obj_body)
    center = data.xpos[oid].copy()
    size, gtype = None, None
    for gi in range(model.ngeom):
        if model.geom_bodyid[gi] == oid:
            size = model.geom_size[gi].copy()
            gtype = int(model.geom_type[gi])
            break
    # semi-axes by geom type (MuJoCo packs geom_size differently per primitive)
    s = np.asarray(size, float)
    if gtype == mujoco.mjtGeom.mjGEOM_SPHERE:
        semi = np.array([s[0]] * 3)
    elif gtype in (mujoco.mjtGeom.mjGEOM_CYLINDER, mujoco.mjtGeom.mjGEOM_CAPSULE):
        semi = np.array([s[0], s[0], s[1]])            # radius, radius, half-len
    else:                                               # ellipsoid / box: 3 semi-axes
        semi = s[:3]
    palm = data.xpos[_bid(model, "rh_palm")].copy()
    approach = center - palm
    distance = float(np.linalg.norm(approach))
    unit = approach / (distance + 1e-9)
    radii = np.sort(semi)                              # ascending semi-axes
    width = float(2 * radii[0])                        # tightest cross-section
    span = float(2 * radii[-1])                        # longest extent
    tips = {k: data.xpos[_bid(model, b)].copy() for k, b in TIP_BODY.items()}
    return {
        "object": obj_body,
        "center": [round(float(x), 3) for x in center],
        "size": [round(float(x), 3) for x in semi],     # semi-axes (a, b, c)
        "gtype": gtype,
        "palm": [round(float(x), 3) for x in palm],
        "approach": [round(float(x), 3) for x in unit],
        "distance": round(distance, 3),
        "width": round(width, 3),
        "span": round(span, 3),
        "_center": center, "_tips": tips,        # raw vectors for the planner
    }


def plan_grasp(model, data, percept):
    """Decide an opposition grasp sized to the perceived object."""
    w = percept["width"]
    r = w / 2.0
    # smaller objects need a tighter curl; clamp to a sane band
    strength = float(np.clip(1.15 - 3.0 * r, 0.5, 1.0))
    aperture = round(w + 0.02, 3)                       # pre-grasp opening
    spread = float(np.clip(percept["span"] * 1.2, 0.0, 0.349))
    # Wrist aim: ideal tilt toward the object's vertical offset, but only a gentle
    # fraction is commanded — a large tilt swings the fingers out of the grasp
    # envelope and knocks a free object off. The object sits in the neutral
    # envelope, so we keep the wrist near neutral while still reporting the aim.
    wrist_ideal = float(np.clip(percept["approach"][2] * 0.8, -0.698, 0.489))
    wrist_aim = float(np.clip(wrist_ideal * 0.15, -0.2, 0.2))

    preshape = {"rh_A_WRJ1": wrist_aim, "rh_A_THJ4": 0.6}
    closed = {"rh_A_WRJ1": wrist_aim}
    for i, f in enumerate(FINGERS):
        preshape[f"rh_A_{f}J3"] = 0.0
        preshape[f"rh_A_{f}J0"] = 0.0
        closed[f"rh_A_{f}J3"] = round(strength * _J3_CLOSED, 3)
        closed[f"rh_A_{f}J0"] = round(strength * _J0_CLOSED, 3)
    # thumb opposition + flex (the opposing contact)
    closed["rh_A_THJ4"] = 1.1
    closed["rh_A_THJ2"] = 0.5
    closed["rh_A_THJ1"] = 1.0

    # An opposition grasp opposes the thumb against the four fingers by design;
    # how good the actual hold is is judged from the achieved contacts (see
    # grasp_quality / snapshot), a cheap stand-in for the paper's force-closure
    # test. At plan time we report the size-fit (how well the object fits the
    # graspable band) as the predicted quality.
    fit = float(np.clip(1.0 - abs(w - 0.045) / 0.045, 0.0, 1.0))

    return {
        "type": "opposition",
        "strength": round(strength, 3),
        "aperture": aperture,
        "spread": round(spread, 3),
        "wrist_aim": round(wrist_aim, 3),
        "contacts": 5,                                   # planned finger count
        "fit": round(fit, 3),
        "opposition": 1.0,                               # thumb-vs-fingers by design
        "quality": round(0.5 + 0.5 * fit, 3),            # predicted; refined live
        "preshape": preshape,
        "closed": closed,
    }


def grasp_quality(touched):
    """Force-closure-style score from the achieved contacts: coverage + whether
    the thumb participates (opposition). Cheap proxy for the paper's test."""
    n = len(touched)
    opposition = 1.0 if "TH" in touched else (0.5 if n else 0.0)
    quality = round(min(1.0, 0.4 + 0.12 * n), 3)
    return opposition, quality


def snapshot(percept, plan, phase=None, touched=None):
    """JSON-safe view of the whole grasp plan for /status and the OS terminal."""
    p = {k: v for k, v in (percept or {}).items() if not k.startswith("_")}
    out = {"perception": p, "plan": None}
    if plan is not None:
        out["plan"] = {k: v for k, v in plan.items()
                       if k not in ("preshape", "closed")}
    if phase is not None:
        out["phase"] = phase
    if touched is not None:
        out["touched"] = sorted(touched)
        out["n_contact"] = len(touched)
        # refine opposition/quality from what the hand actually achieved
        if out["plan"] is not None and touched:
            opp, q = grasp_quality(touched)
            out["plan"] = {**out["plan"], "opposition": opp, "quality": q}
    return out