"""Cerebral vessel segmentation and trajectory clearance-to-vessel.

Hitting a blood vessel is the dominant safety concern when placing a depth
electrode or biopsy needle, so a planner must segment the vasculature and report
how close a trajectory comes to it. This module:

* extracts vessels from an angiographic volume (CTA / MR angiography / TOF) with
  a **multiscale Hessian vesselness filter** (Frangi/Sato) — no learned weights,
  no network, and only the existing scikit-image dependency; vessels are bright
  on these sequences, so the default enhances bright tubular structures;
* thresholds the vesselness response into a binary vessel mask; and
* computes the **minimum clearance (mm) from a trajectory to the vessel tree**
  with a KD-tree (fast even for large masks), optionally against the
  *predicted deformed* vasculature — deformation-aware vessel avoidance, the
  same differentiator applied to vessels.

A pretrained deep-learning backend can be slotted in behind ``method="model"``
later; the Frangi backend is the always-available default.

IMPORTANT: Research software only. Not a medical device. Automated vessel
segmentations are unverified and must never be relied on for patient care.
See README and docs/INTENDED_USE.md.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band


def vesselness(volume: Volume, scales=(1.0, 2.0, 3.0, 4.0), bright_vessels=True) -> Volume:
    """Multiscale Frangi vesselness response (same grid as `volume`).

    scales          : Hessian sigmas (mm-ish, in voxels) spanning vessel radii.
    bright_vessels  : True for CTA/MRA/TOF (vessels bright); False for dark vessels.
    """
    from skimage.filters import frangi
    data = np.asarray(volume.data, dtype=np.float64)
    resp = frangi(data, sigmas=list(scales), black_ridges=not bright_vessels)
    return Volume(np.nan_to_num(resp).astype(np.float32), volume.affine.copy())


def vessel_mask(volume: Volume, scales=(1.0, 2.0, 3.0, 4.0), bright_vessels=True,
                threshold=None, percentile=99.0) -> Volume:
    """Binary vessel mask from thresholded vesselness.

    threshold : absolute cut on the vesselness response; if None, use the given
                percentile of the positive response (auto, data-adaptive).
    """
    resp = vesselness(volume, scales, bright_vessels)
    r = np.asarray(resp.data)
    if threshold is None:
        pos = r[r > 0]
        threshold = float(np.percentile(pos, percentile)) if pos.size else np.inf
    return Volume((r >= threshold).astype(np.uint8), volume.affine.copy())


def segment_vessels(volume: Volume, method="frangi", **kwargs) -> Volume:
    """Segment vasculature into a binary mask Volume.

    method="frangi" (default) uses the multiscale vesselness filter. A learned
    backend ("model") is left as an extension point.
    """
    if method == "frangi":
        return vessel_mask(volume, **kwargs)
    if method == "model":
        raise NotImplementedError(
            "A pretrained vessel-segmentation backend is not wired in yet; "
            'use method="frangi" (multiscale vesselness).')
    raise ValueError(f"Unknown vessel segmentation method {method!r}.")


def _mask_world_points(mask: Volume, deformation=None) -> np.ndarray:
    """World (mm) coordinates of the mask's set voxels, optionally deformed."""
    idx = np.argwhere(np.asarray(mask.data) > 0)
    if idx.size == 0:
        return np.empty((0, 3))
    homog = np.hstack([idx.astype(np.float64), np.ones((idx.shape[0], 1))])
    world = (mask.affine @ homog.T).T[:, :3]
    if deformation is not None:
        world = world + deformation.field(world)
    return world


def clearance_to_mask(traj: Trajectory, mask: Volume, n=256, deformation=None):
    """Minimum distance (mm) from the trajectory to any set voxel of `mask`.

    Uses a KD-tree over the mask points, so it stays fast even for a dense vessel
    tree. Returns (distance_mm, world_point_on_traj). If `deformation` is given,
    the mask points are displaced first (deformation-aware clearance).
    """
    from scipy.spatial import cKDTree
    world = _mask_world_points(mask, deformation)
    if world.shape[0] == 0:
        return float("inf"), None
    tree = cKDTree(world)
    pts = traj.sample(n)
    d, _ = tree.query(pts)
    j = int(np.argmin(d))
    return float(d[j]), pts[j]


def vessel_clearance(traj: Trajectory, vessel_mask_vol: Volume, deformation=None,
                     n=256) -> dict:
    """Trajectory clearance to the vessel tree, static and (optionally) deformed.

    Returns clearance (mm) and a risk band; when `deformation` is supplied, adds
    the deformation-aware clearance and the change vs. static.
    """
    d, p = clearance_to_mask(traj, vessel_mask_vol, n=n)
    out = {"vessel_clearance_mm": round(d, 2) if np.isfinite(d) else float("inf"),
           "vessel_band": risk_band(d), "nearest_point": p}
    if deformation is not None:
        dd, _ = clearance_to_mask(traj, vessel_mask_vol, n=n, deformation=deformation)
        out.update({"deformed_vessel_clearance_mm": round(dd, 2) if np.isfinite(dd) else float("inf"),
                    "deformed_vessel_band": risk_band(dd),
                    "vessel_clearance_change_mm": round(dd - d, 2) if np.isfinite(dd) and np.isfinite(d) else None})
    return out


def vessel_risk(traj: Trajectory, volume: Volume, deformation=None,
                scales=(1.0, 2.0, 3.0, 4.0), bright_vessels=True,
                threshold=None, percentile=99.0, n=256) -> dict:
    """One-call convenience: segment vessels from an angiographic volume, then
    report trajectory clearance-to-vessel (static and deformation-aware)."""
    mask = vessel_mask(volume, scales, bright_vessels, threshold, percentile)
    out = vessel_clearance(traj, mask, deformation=deformation, n=n)
    out["vessel_voxels"] = int(np.asarray(mask.data).sum())
    return out
