Presa-Mano/src/presa_mano/api/grasp_routes.py

"""Grasp planning + execution endpoints."""

from __future__ import annotations

from typing import Any, Optional

import numpy as np
import time
from dataclasses import asdict
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from ..calibration import CalibrationStore
from ..clients import G1Unavailable
from ..config import get_config
from ..geometry import transform_points
from ..grasp import GraspPlanner, compute_axes
from ..grasp.visual_servo import (
    ServoError,
    WorkPlane,
    fit_plane_svd,
    get_work_plane,
    plan as servo_plan,
    set_work_plane,
)
from ..perception import ObjectDetector, filter_outliers, nuwa_to_cv_frame
from .perception_routes import _bbox_to_roi
from ..timing import Stopwatch
from . import _deps

router = APIRouter()


class PlanReq(BaseModel):
    conf: Optional[float] = None
    class_name: Optional[str] = None


class ExecuteReq(PlanReq):
    dry_run: bool = True


def _object_points_in_base(conf: float | None, class_name: str | None, sw: Stopwatch
                           ) -> tuple[np.ndarray, dict[str, Any]]:
    cfg = get_config()
    nuwa = _deps.nuwa()
    yolo = _deps.yolo()

    side = cfg.get("robot", {}).get("wrist_side", "right")
    T_base_cam = CalibrationStore(_deps.calibration_file_path(side)).T_base_camera()
    if T_base_cam is None:
        raise HTTPException(status_code=412, detail=f"calibration missing for side='{side}': run /api/calibration/solve first")

    detector = ObjectDetector(yolo, nuwa)
    try:
        with sw.step("yolo"):
            dets = detector.detect(conf=conf)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"YOLO detect failed: {e}")
    if class_name:
        dets = [d for d in dets if (d.get("class") or "").lower() == class_name.lower()]
    if not dets:
        raise HTTPException(status_code=404, detail="no object detected")
    dets.sort(key=lambda d: d["confidence"], reverse=True)
    top = dets[0]
    bbox = _bbox_to_roi(top["bbox"])

    pts_base = _points_via_bbox(nuwa, T_base_cam, bbox, sw)
    meta = {
        "detection": top,
        "T_base_camera": T_base_cam.tolist(),
        "n_points": int(len(pts_base)),
        "source": "bbox_roi",
    }
    return pts_base, meta


def _points_via_bbox(nuwa, T_base_cam: np.ndarray, bbox: tuple[int, int, int, int],
                     sw: Stopwatch) -> np.ndarray:
    """Legacy path: bbox → Nuwa ROI cloud → filter outliers → base frame.

    Raises HTTPException(422) on <10 valid depth points (existing behaviour).
    """
    cfg = get_config()
    u_min, v_min, u_max, v_max = bbox
    try:
        with sw.step("pointcloud_roi"):
            pc, _meta = nuwa.pointcloud_roi(
                u_min, v_min, u_max, v_max,
                step=int(cfg.get("grasp", {}).get("pointcloud_step", 1)),
                max_depth_mm=int(cfg.get("grasp", {}).get("pointcloud_max_depth_mm", 2000)),
            )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Nuwa ROI pointcloud failed: {e}")
    with sw.step("filter_outliers"):
        cropped = filter_outliers(pc)
    if len(cropped) < 10:
        raise HTTPException(status_code=422, detail=f"too few valid depth points in bbox: {len(cropped)}")
    with sw.step("to_base_frame"):
        pts_cv = nuwa_to_cv_frame(cropped[:, :3])
        pts_base = transform_points(T_base_cam, pts_cv)
    return pts_base


@router.post("/plan")
def plan(req: PlanReq) -> dict[str, Any]:
    sw = Stopwatch()
    cfg = get_config()
    pts_base, meta = _object_points_in_base(req.conf, req.class_name, sw)
    with sw.step("pca_axes"):
        axes = compute_axes(pts_base)
    planner = GraspPlanner(
        approach_offset_m=float(cfg["grasp"]["approach_offset_m"]),
        max_object_width_m=float(cfg["grasp"]["max_object_width_m"]),
    )
    with sw.step("plan"):
        g = planner.plan(axes)
    return {"plan": g.to_dict(), "timings_ms": sw.as_dict(), **meta}


@router.post("/execute")
def execute(req: ExecuteReq) -> dict[str, Any]:
    sw = Stopwatch()
    cfg = get_config()
    side = cfg["robot"]["wrist_side"]
    pts_base, meta = _object_points_in_base(req.conf, req.class_name, sw)
    with sw.step("pca_axes"):
        axes = compute_axes(pts_base)
    planner = GraspPlanner(
        approach_offset_m=float(cfg["grasp"]["approach_offset_m"]),
        max_object_width_m=float(cfg["grasp"]["max_object_width_m"]),
    )
    with sw.step("plan"):
        g = planner.plan(axes)

    open_action = "RIGHT_OPEN" if side == "right" else "LEFT_OPEN"
    close_action = "RIGHT_CLOSE" if side == "right" else "LEFT_CLOSE"

    g1 = _deps.g1(dry_run=req.dry_run)
    steps: list[dict[str, Any]] = []
    try:
        with sw.step("hand_open"):
            steps.append({"step": "hand_open", "result": g1.hand_action(open_action)})
        with sw.step("pre_grasp"):
            steps.append({"step": "pre_grasp", "result": g1.wrist_absolute(
                side=side,
                position=g.pre_grasp_position.tolist(),
                quaternion=g.grasp_quaternion.tolist(),
            )})
        with sw.step("grasp_approach"):
            steps.append({"step": "grasp_approach", "result": g1.wrist_absolute(
                side=side,
                position=g.grasp_position.tolist(),
                quaternion=g.grasp_quaternion.tolist(),
            )})
        with sw.step("hand_close"):
            steps.append({"step": "hand_close", "result": g1.hand_action(close_action)})
    except G1Unavailable as e:
        raise HTTPException(status_code=502, detail=f"G1 unavailable: {e}")
    return {"plan": g.to_dict(), "steps": steps, "timings_ms": sw.as_dict(), **meta}


# ---------------------------------------------------------------------------
# Visual servo (marker-relative, no static hand-eye calibration).
# ---------------------------------------------------------------------------

class VisualServoReq(BaseModel):
    conf: Optional[float] = None
    class_name: Optional[str] = None
    side: Optional[str] = None
    min_distance_m: Optional[float] = None   # default from config (0.1 m)
    step_fraction: Optional[float] = None    # default from config (0.5)
    # Override for grasp.visual_servo.wrist_marker_rvec (radians). Useful for
    # live tuning from the UI without persisting to config on every try.
    wrist_marker_rvec: Optional[list[float]] = None


class VisualServoStepReq(VisualServoReq):
    dry_run: bool = True
    timeout_s: float = 5.0


def _servo_plan_from_req(req: VisualServoReq, sw: Stopwatch):
    cfg = get_config()
    vs_cfg = (cfg.get("grasp", {}) or {}).get("visual_servo", {}) or {}
    side = req.side or cfg["robot"]["wrist_side"]
    min_dist = float(req.min_distance_m
                     if req.min_distance_m is not None
                     else vs_cfg.get("min_distance_m", 0.10))
    step_frac = float(req.step_fraction
                      if req.step_fraction is not None
                      else vs_cfg.get("step_fraction", 0.5))
    # Request-level rvec override: thread it into a shallow-copied cfg so the
    # core plan() function keeps reading from a single place.
    if req.wrist_marker_rvec is not None:
        cfg = {**cfg, "grasp": {**(cfg.get("grasp") or {}),
                                "visual_servo": {**vs_cfg,
                                                 "wrist_marker_rvec": list(req.wrist_marker_rvec)}}}
    marker_id = int(cfg["aruco"]["wrist_marker"]["id"])
    size_mm = float(cfg["aruco"]["wrist_marker"]["size_mm"])
    pc_step = int(cfg.get("grasp", {}).get("pointcloud_step", 1))
    pc_max_depth = int(cfg.get("grasp", {}).get("pointcloud_max_depth_mm", 2000))
    try:
        with sw.step("servo_plan"):
            result = servo_plan(
                side=side, marker_id=marker_id, size_mm=size_mm,
                conf=req.conf, class_name=req.class_name,
                min_distance_m=min_dist, step_fraction=step_frac,
                pointcloud_step=pc_step, pointcloud_max_depth_mm=pc_max_depth,
                cfg=cfg,
                nuwa=_deps.nuwa(), marker=_deps.marker(),
                yolo=_deps.yolo(), g1=_deps.g1(dry_run=False),
            )
    except ServoError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"servo plan failed: {e}")
    return side, min_dist, step_frac, result


@router.post("/visual_servo/plan")
def visual_servo_plan(req: VisualServoReq) -> dict[str, Any]:
    """Perception + next-step computation, no motion. Useful for UI preview."""
    sw = Stopwatch()
    side, min_dist, step_frac, p = _servo_plan_from_req(req, sw)
    return {
        "side": side,
        "min_distance_m": min_dist,
        "step_fraction": step_frac,
        "plan": p.__dict__,
        "timings_ms": sw.as_dict(),
    }


@router.post("/visual_servo/step")
def visual_servo_step(req: VisualServoStepReq) -> dict[str, Any]:
    """Plan + execute one half-step toward the object. Returns the move result
    plus the plan (before the move) so the UI can visualise what just happened;
    the next click re-plans from the new wrist pose."""
    sw = Stopwatch()
    side, min_dist, step_frac, p = _servo_plan_from_req(req, sw)
    move_result: dict[str, Any] | None = None
    if p.done or p.target_position_base is None:
        return {
            "side": side,
            "min_distance_m": min_dist,
            "step_fraction": step_frac,
            "plan": p.__dict__,
            "moved": False,
            "reason": p.reason or "reached",
            "timings_ms": sw.as_dict(),
        }
    g1 = _deps.g1(dry_run=req.dry_run)
    try:
        with sw.step("wrist_absolute"):
            move_result = g1.wrist_absolute(
                side=side,
                position=list(p.target_position_base),
                quaternion=list(p.target_quaternion_base or p.wrist_quaternion_base),
                wait=True,
                timeout_s=req.timeout_s,
            )
    except G1Unavailable as e:
        raise HTTPException(status_code=502, detail=f"G1 unavailable: {e}")
    return {
        "side": side,
        "min_distance_m": min_dist,
        "step_fraction": step_frac,
        "plan": p.__dict__,
        "moved": True,
        "move_result": move_result,
        "timings_ms": sw.as_dict(),
    }


# ---------------------------------------------------------------------------
# Work plane (for the visual-servo tab): fit once from a user-selected ROI,
# stored in-memory (singleton). Subsequent visual_servo/plan calls enrich the
# response with plane projections of marker and object.
# ---------------------------------------------------------------------------

class PlaneFitReq(BaseModel):
    u_min: int = Field(..., ge=0)
    v_min: int = Field(..., ge=0)
    u_max: int = Field(..., gt=0)
    v_max: int = Field(..., gt=0)
    max_depth_mm: Optional[int] = None


@router.post("/visual_servo/plane/fit")
def visual_servo_plane_fit(req: PlaneFitReq) -> dict[str, Any]:
    cfg = get_config()
    nuwa = _deps.nuwa()
    pc_max = req.max_depth_mm or int(cfg.get("grasp", {}).get("pointcloud_max_depth_mm", 2000))
    try:
        pc, _meta = nuwa.pointcloud_roi(
            req.u_min, req.v_min, req.u_max, req.v_max,
            step=1, max_depth_mm=pc_max,
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Nuwa pointcloud_roi failed: {e}")
    if len(pc) < 10:
        raise HTTPException(status_code=422, detail=f"too few depth points in ROI: {len(pc)}")
    pts_cv = nuwa_to_cv_frame(pc[:, :3])
    # Drop gross outliers before SVD (stray far-depth pixels around the edges).
    pts_cv = filter_outliers(pts_cv, n_std=2.0)
    try:
        normal, offset, centroid, rms_mm = fit_plane_svd(pts_cv)
    except ServoError as e:
        raise HTTPException(status_code=422, detail=str(e))
    wp = WorkPlane(
        normal=[float(x) for x in normal],
        offset=float(offset),
        centroid=[float(x) for x in centroid],
        n_inliers=int(len(pts_cv)),
        fit_rms_mm=float(rms_mm),
        roi=[int(req.u_min), int(req.v_min), int(req.u_max), int(req.v_max)],
        created_at=time.time(),
    )
    set_work_plane(wp)
    return {"ok": True, "plane": asdict(wp)}


@router.get("/visual_servo/plane")
def visual_servo_plane_get() -> dict[str, Any]:
    wp = get_work_plane()
    return {"set": wp is not None, "plane": asdict(wp) if wp is not None else None}


@router.delete("/visual_servo/plane")
def visual_servo_plane_clear() -> dict[str, Any]:
    set_work_plane(None)
    return {"ok": True, "set": False}