"""Ventricular catheter planning — external ventricular drain (EVD) / shunt.

Placing a catheter into the frontal horn of a lateral ventricle (for CSF
drainage in hydrocephalus/IVH, or a shunt proximal catheter) is one of the most
common neurosurgical procedures. The target is the frontal horn; the corridor
should reach the ventricle with a safe length. This module segments the
ventricles (CSF), locates the frontal-horn target, and plans/scored the catheter.

Pure NumPy + scipy; unit-tested headlessly. Not a medical device.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band
from .hematoma import traversal_length
from .optimizer import candidate_entries
from .vessels import _mask_world_points
from .session import default_entry_radius


def segment_ventricles(volume: Volume, hu_lo=-5.0, hu_hi=18.0, min_volume_cc=0.5) -> Volume:
    """Ventricle (CSF) mask from CT — CSF HU inside the brain, keeping the large
    central components (drops thin sulcal CSF at the surface)."""
    from scipy import ndimage
    from .ct import csf_mask
    csf = np.asarray(csf_mask(volume, hu_lo, hu_hi).data)
    lab, n = ndimage.label(csf)
    if not n:
        return Volume(np.zeros(volume.shape, np.uint8), volume.affine.copy())
    sizes = np.bincount(lab.ravel()); sizes[0] = 0
    voxel_cc = float(np.prod(volume.spacing)) / 1000.0
    keep = [i for i in range(1, n + 1) if sizes[i] * voxel_cc >= min_volume_cc]
    m = np.isin(lab, keep) if keep else (lab == int(np.argmax(sizes)))
    return Volume(m.astype(np.uint8), volume.affine.copy())


def frontal_horn_target(ventricle_mask: Volume, side="right") -> np.ndarray:
    """Target point in the frontal horn: the anterior portion of the ventricle
    on the chosen side (anterior = +y in RAS; right = +x)."""
    idx = np.argwhere(np.asarray(ventricle_mask.data) > 0)
    if idx.size == 0:
        return None
    homog = np.hstack([idx.astype(np.float64), np.ones((idx.shape[0], 1))])
    world = (ventricle_mask.affine @ homog.T).T[:, :3]
    y = world[:, 1]
    front = world[y >= np.percentile(y, 70)]                 # anterior third
    if front.shape[0] == 0:
        front = world
    xm = np.median(front[:, 0])
    sub = front[front[:, 0] >= xm] if side == "right" else front[front[:, 0] < xm]
    return (sub.mean(axis=0) if sub.shape[0] else front.mean(axis=0))


def plan_ventricular_catheter(volume: Volume, ventricle_mask: Volume, side="right",
                              entry=None, structures=None, axis=(0, 0, 1),
                              max_angle_deg=45.0, n_theta=6, n_phi=18, n=200):
    """Plan a catheter to the frontal horn. With `entry` given, score that path;
    otherwise search frontal-approach corridors and pick the one that best
    cannulates the ventricle while clearing any supplied structures.

    Returns a dict with the trajectory, catheter length, whether it reaches the
    ventricle, ventricle traversal, and clearance.
    """
    from scipy.spatial import cKDTree
    target = frontal_horn_target(ventricle_mask, side)
    if target is None:
        return None
    structures = structures or []
    trees = {s.name: (cKDTree(_mask_world_points(s.mask)) if s.mask.data.any() else None)
             for s in structures}

    def score(tr):
        trav = traversal_length(tr, ventricle_mask, n)
        clears = {s.name: (float(np.min(trees[s.name].query(tr.sample(n))[0]))
                           if trees[s.name] is not None else float("inf")) for s in structures}
        minc = min(clears.values()) if clears else float("inf")
        tipv = np.round(ventricle_mask.world_to_voxel(tr.target)).astype(int)
        reaches = bool(ventricle_mask.in_bounds(tipv) and ventricle_mask.data[tuple(tipv)] > 0)
        return trav, clears, minc, reaches

    if entry is not None:
        tr = Trajectory(np.asarray(entry, float), target)
        trav, clears, minc, reaches = score(tr)
        best = (tr, trav, clears, minc, reaches)
    else:
        radius = default_entry_radius(volume, target, axis)
        best = None
        for e in candidate_entries(target, radius, axis=axis, max_angle_deg=max_angle_deg,
                                   n_theta=n_theta, n_phi=n_phi):
            tr = Trajectory(e, target)
            trav, clears, minc, reaches = score(tr)
            feasible = all(clears[s.name] >= s.min_clearance_mm for s in structures)
            key = (reaches, feasible, trav, minc)
            if best is None or key > best[5]:
                best = (tr, trav, clears, minc, reaches, key)
        best = best[:5]

    tr, trav, clears, minc, reaches = best
    return {"trajectory": tr, "entry": tr.entry, "target": target,
            "catheter_length_mm": round(tr.length, 1), "reaches_ventricle": reaches,
            "ventricle_traversal_mm": round(trav, 1),
            "min_clearance_mm": round(minc, 2) if np.isfinite(minc) else float("inf"),
            "band": risk_band(minc) if structures else None,
            "clearances": {k: round(v, 2) for k, v in clears.items()}}
