"""Trajectory planning primitives and geometric safety metrics (world/mm space)."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .volume import Volume


@dataclass
class Trajectory:
    """A straight-line stereotactic trajectory in world (mm) coordinates."""
    entry: np.ndarray
    target: np.ndarray

    def __post_init__(self):
        self.entry = np.asarray(self.entry, dtype=np.float64)
        self.target = np.asarray(self.target, dtype=np.float64)

    @property
    def vector(self) -> np.ndarray:
        return self.target - self.entry

    @property
    def length(self) -> float:
        """Trajectory length / target depth from entry (mm)."""
        return float(np.linalg.norm(self.vector))

    @property
    def direction(self) -> np.ndarray:
        L = self.length
        return self.vector / L if L > 0 else self.vector

    def angle_from_axis(self, axis=(0, 0, 1)) -> float:
        """Angle (deg) between trajectory and a world axis (default +Z / superior)."""
        a = np.asarray(axis, dtype=np.float64)
        a = a / np.linalg.norm(a)
        c = abs(float(np.dot(self.direction, a)))
        return float(np.degrees(np.arccos(min(1.0, c))))

    def sample(self, n=128) -> np.ndarray:
        """n points along the trajectory, shape (n,3) in world mm."""
        t = np.linspace(0.0, 1.0, n)[:, None]
        return self.entry[None, :] * (1 - t) + self.target[None, :] * t

    # ---- safety metrics vs. a labelled structure mask ----
    def min_clearance_to_mask(self, mask: Volume, n=256):
        """Minimum distance (mm) from the trajectory to any voxel of a binary
        structure `mask`. Returns (distance_mm, world_point_on_traj).
        Distance is >0 when clear; ~0 (or negative not modelled) when intersecting.
        """
        pts = self.sample(n)
        idx = np.argwhere(np.asarray(mask.data) > 0)
        if idx.size == 0:
            return float("inf"), None
        # structure voxel centres in world mm
        ones = np.ones((idx.shape[0], 1))
        homog = np.hstack([idx.astype(np.float64), ones])
        world = (mask.affine @ homog.T).T[:, :3]
        best_d, best_p = float("inf"), None
        for p in pts:
            d = np.min(np.linalg.norm(world - p[None, :], axis=1))
            if d < best_d:
                best_d, best_p = float(d), p
        return best_d, best_p

    def passes_through(self, mask: Volume, n=256) -> bool:
        """True if any sampled point falls inside a binary mask voxel."""
        for p in self.sample(n):
            v = np.round(mask.world_to_voxel(p)).astype(int)
            if mask.in_bounds(v) and mask.data[tuple(v)] > 0:
                return True
        return False


def risk_band(clearance_mm: float) -> str:
    """Simple geometric risk band from min clearance to critical structures."""
    if clearance_mm < 3.0:
        return "HIGH"
    if clearance_mm < 6.0:
        return "CAUTION"
    return "LOW"
