# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""A cuRobo-free PRM (Probabilistic Roadmap) motion planner in end-effector task space.

This is the motion-planning primitive for the Galbot dual-parallel TAMP integration
when cuRobo cannot be compiled. It plans a collision-free polyline for the end
effector between waypoints, checking straight-line edges against axis-aligned box
obstacles (oven body, table, fridge, ...) inflated by an EEF clearance radius.

The env's RMPFlow controller turns each EEF waypoint into joint commands, so a
task-space roadmap over reachable EEF positions is sufficient to drive the arm
around obstacles; grasp orientations come from :mod:`mesh_grasp`.

Pure numpy + scipy (KD-tree). Deterministic given ``seed``.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np


@dataclass
class Box:
    """Axis-aligned box obstacle in world/env-local frame."""

    center: np.ndarray
    half: np.ndarray  # half extents (x,y,z)

    def contains(self, p: np.ndarray, pad: float = 0.0) -> bool:
        d = np.abs(np.asarray(p) - self.center)
        return bool(np.all(d <= (self.half + pad)))

    def segment_hits(self, a: np.ndarray, b: np.ndarray, pad: float = 0.0, steps: int = 24) -> bool:
        """Conservative segment-box test by dense sampling (robust, simple)."""
        for t in np.linspace(0.0, 1.0, steps):
            if self.contains(a + t * (b - a), pad):
                return True
        return False


@dataclass
class PRM:
    bounds_lo: np.ndarray
    bounds_hi: np.ndarray
    obstacles: list[Box] = field(default_factory=list)
    clearance: float = 0.05      # EEF sphere radius for collision inflation
    num_samples: int = 300
    k: int = 10
    seed: int = 0

    def _free_point(self, p: np.ndarray) -> bool:
        if np.any(p < self.bounds_lo) or np.any(p > self.bounds_hi):
            return False
        return not any(o.contains(p, self.clearance) for o in self.obstacles)

    def _free_edge(self, a: np.ndarray, b: np.ndarray) -> bool:
        return not any(o.segment_hits(a, b, self.clearance) for o in self.obstacles)

    def plan(self, start: np.ndarray, goal: np.ndarray) -> np.ndarray | None:
        """Return an (M,3) collision-free polyline from start to goal, or None."""
        from scipy.spatial import cKDTree

        start = np.asarray(start, dtype=np.float64)
        goal = np.asarray(goal, dtype=np.float64)
        rng = np.random.default_rng(self.seed)

        # direct connection first
        if self._free_point(start) and self._free_point(goal) and self._free_edge(start, goal):
            return np.stack([start, goal]).astype(np.float32)

        # sample free nodes
        nodes = [start, goal]
        tries = 0
        while len(nodes) < self.num_samples + 2 and tries < self.num_samples * 30:
            tries += 1
            p = rng.uniform(self.bounds_lo, self.bounds_hi)
            if self._free_point(p):
                nodes.append(p)
        nodes = np.array(nodes)
        n = len(nodes)
        if n < 3:
            return None

        # k-NN graph with collision-free edges
        tree = cKDTree(nodes)
        adj: dict[int, list[tuple[int, float]]] = {i: [] for i in range(n)}
        for i in range(n):
            dists, idxs = tree.query(nodes[i], k=min(self.k + 1, n))
            for d, j in zip(np.atleast_1d(dists), np.atleast_1d(idxs)):
                j = int(j)
                if j == i:
                    continue
                if self._free_edge(nodes[i], nodes[j]):
                    adj[i].append((j, float(d)))

        # Dijkstra start(0) -> goal(1)
        import heapq

        dist = {0: 0.0}
        prev: dict[int, int] = {}
        pq = [(0.0, 0)]
        while pq:
            d, u = heapq.heappop(pq)
            if u == 1:
                break
            if d > dist.get(u, np.inf):
                continue
            for v, w in adj[u]:
                nd = d + w
                if nd < dist.get(v, np.inf):
                    dist[v] = nd
                    prev[v] = u
                    heapq.heappush(pq, (nd, v))
        if 1 not in dist:
            return None
        # backtrace
        path = [1]
        while path[-1] != 0:
            path.append(prev[path[-1]])
        path = path[::-1]
        return nodes[path].astype(np.float32)


def shortcut(path: np.ndarray, obstacles: list[Box], clearance: float, iters: int = 100,
             seed: int = 0) -> np.ndarray:
    """Randomized path shortcutting: greedily replace subpaths with straight edges."""
    if path is None or len(path) < 3:
        return path
    rng = np.random.default_rng(seed)
    pts = [p for p in path]
    for _ in range(iters):
        if len(pts) < 3:
            break
        i = rng.integers(0, len(pts) - 2)
        j = rng.integers(i + 2, len(pts))
        a, b = pts[i], pts[j]
        if not any(o.segment_hits(a, b, clearance) for o in obstacles):
            pts = pts[: i + 1] + pts[j:]
    return np.array(pts, dtype=np.float32)


def resample_polyline(path: np.ndarray, n: int) -> np.ndarray:
    """Arc-length resample an (M,3) polyline to exactly n points."""
    path = np.asarray(path, dtype=np.float64)
    seg = np.linalg.norm(np.diff(path, axis=0), axis=1)
    cum = np.concatenate([[0.0], np.cumsum(seg)])
    total = cum[-1] if cum[-1] > 1e-9 else 1.0
    ts = np.linspace(0.0, total, n)
    out = np.empty((n, 3))
    for d in range(3):
        out[:, d] = np.interp(ts, cum, path[:, d])
    return out.astype(np.float32)