"""Task-space path planning: a thin wrapper over the PRM in scripts/yam_prm.py.

Kept separate from the arm controller so a task can swap in a different planner (straight line,
cuRobo, a learned policy) without touching the tracking/grasp code.
"""
from __future__ import annotations

import numpy as np

# planning-only blockers behind the robot: never rendered, only used to keep the planned
# end-effector path from sweeping backwards through the arm's own base
DEFAULT_WALLS = [
    (np.array([-0.38, -0.15, 0.70]), np.array([0.02, 0.45, 0.28])),
    (np.array([-0.10, -0.45, 0.70]), np.array([0.45, 0.02, 0.28])),
]


def _prm_module():
    import importlib, os, sys
    scripts = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(
        os.path.dirname(os.path.abspath(__file__))))), "scripts")
    if scripts not in sys.path:
        sys.path.insert(0, scripts)
    return importlib.import_module("yam_prm")


def plan_path(start, goal, arm_root, walls=None, clearance=0.05, samples=300, resample=12):
    """Collision-free end-effector polyline from `start` to `goal`, both in the arm's root frame.

    Falls back to a straight line if the roadmap finds nothing, so a task never dies here -- a
    blocked path shows up as a tracking error the task can see, not an exception.
    """
    yam_prm = _prm_module()
    obstacles = [yam_prm.Box(center=(c-np.asarray(arm_root, np.float64)), half=h)
                 for c, h in (walls if walls is not None else DEFAULT_WALLS)]
    prm = yam_prm.PRM(bounds_lo=np.array([-0.05, -0.35, -0.03]),
                      bounds_hi=np.array([0.55, 0.45, 0.40]),
                      obstacles=obstacles, clearance=clearance, num_samples=samples, k=12, seed=1)
    path = prm.plan(np.asarray(start, np.float64), np.asarray(goal, np.float64))
    if path is None:
        return [np.asarray(goal, np.float32)]
    short = yam_prm.shortcut(path, obstacles, clearance, iters=200, seed=2)
    dense = yam_prm.resample_polyline(short, resample).astype(np.float32)
    return list(dense[1:])          # drop the start point; the executor begins where it is