"""Intracerebral-haemorrhage (ICH) evacuation planning.

Minimally invasive corridor evacuation of a supratentorial clot is now
RCT-supported (ENRICH, NEJM 2024), and automated pipelines segment the
haematoma and plan a puncture trajectory that maximises clot traversal while
avoiding eloquent structures and vessels. This module provides that pathway on
top of Sulcora's optimizer:

* ``segment_hematoma`` — a classical CT segmentation (acute blood is hyperdense,
  ~45-95 HU) by thresholding + largest connected component; no learned weights,
  only scipy. A deep-learning backend can be slotted in later.
* ``hematoma_axes`` — PCA centroid + principal axes (the clot's long axis), the
  basis published ICH planners use for targeting.
* ``plan_hematoma`` — searches entry corridors and ranks them by **clot
  traversal length** (so the catheter passes through as much clot as possible),
  subject to hard safety clearances, optionally under predicted deformation.
  Centroid targeting alone is known to sit ~17 mm from surgeons' chosen targets,
  so ranking by traversal (not just hitting the centre) is the refinement.

IMPORTANT: Research software only. Not a medical device. See docs/INTENDED_USE.md.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band
from .optimizer import candidate_entries
from .vessels import _mask_world_points
from .session import default_entry_radius


def segment_hematoma(volume: Volume, hu_lo=45.0, hu_hi=100.0, min_volume_cc=1.0,
                     keep_largest=True, close=True) -> Volume:
    """Segment an acute clot from non-contrast CT by hyperdensity (HU window),
    then keep the largest connected component (or all components >= min volume).
    """
    from scipy import ndimage
    d = np.asarray(volume.data)
    m = (d >= hu_lo) & (d <= hu_hi)
    if close:
        m = ndimage.binary_closing(m, iterations=1)
    lab, n = ndimage.label(m)
    if n == 0:
        return Volume(np.zeros(volume.shape, np.uint8), volume.affine.copy())
    sizes = np.bincount(lab.ravel())
    sizes[0] = 0                                   # background
    voxel_cc = float(np.prod(volume.spacing)) / 1000.0
    if keep_largest:
        keep = {int(np.argmax(sizes))}
    else:
        keep = {i for i in range(1, n + 1) if sizes[i] * voxel_cc >= min_volume_cc}
    mask = np.isin(lab, list(keep)).astype(np.uint8) if keep else np.zeros(volume.shape, np.uint8)
    return Volume(mask, volume.affine.copy())


def _classify_location(centroid, brain_center, brain_half) -> str:
    """Coarse (heuristic, no registration) location of a clot for triage: lobar
    vs deep vs infratentorial — clinically relevant since ENRICH benefit was
    lobar. RAS world coords: +z superior, +x right, +y anterior."""
    rel = (np.asarray(centroid, float) - brain_center) / np.maximum(brain_half, 1e-6)
    if rel[2] < -0.35:
        return "infratentorial (cerebellar/brainstem)"
    if float(np.linalg.norm(rel[:2])) < 0.45:
        return "deep (basal ganglia / thalamic)"
    return "lobar"


def list_hematomas(volume: Volume, hu_lo=45.0, hu_hi=100.0, min_volume_cc=1.0) -> list:
    """Detect and list the haematoma(s) to evacuate — one per connected clot —
    each with world centroid, volume (cc), and a coarse location. This is the
    *target list for ICH*: the places to evacuate, not anatomical structures.
    Sorted largest-first.
    """
    from scipy import ndimage
    from .ct import brain_mask
    d = np.asarray(volume.data)
    m = ndimage.binary_closing((d >= hu_lo) & (d <= hu_hi), iterations=1)
    lab, n = ndimage.label(m)
    if not n:
        return []
    voxel_cc = float(np.prod(volume.spacing)) / 1000.0
    bidx = np.argwhere(np.asarray(brain_mask(volume).data) > 0)
    if bidx.size:
        bw = (volume.affine @ np.hstack([bidx.astype(float), np.ones((bidx.shape[0], 1))]).T).T[:, :3]
        bcenter, bhalf = bw.mean(axis=0), (bw.max(axis=0) - bw.min(axis=0)) / 2.0
    else:
        bcenter, bhalf = np.zeros(3), np.ones(3)
    out = []
    for i in range(1, n + 1):
        vox = np.argwhere(lab == i)
        vol_cc = vox.shape[0] * voxel_cc
        if vol_cc < min_volume_cc:
            continue
        world = (volume.affine @ np.hstack([vox.astype(float), np.ones((vox.shape[0], 1))]).T).T[:, :3]
        centroid = world.mean(axis=0)
        out.append({"centroid": centroid, "volume_cc": round(vol_cc, 1),
                    "location": _classify_location(centroid, bcenter, bhalf),
                    "voxels": int(vox.shape[0])})
    out.sort(key=lambda h: h["volume_cc"], reverse=True)
    return out


def hematoma_axes(mask: Volume) -> dict:
    """Centroid, principal axes (rows; axes[0] = long axis), extents (mm) along
    each axis, and volume (cc) of a segmented clot."""
    idx = np.argwhere(np.asarray(mask.data) > 0)
    if idx.size == 0:
        return {"centroid": None, "axes": None, "lengths": None, "volume_cc": 0.0}
    homog = np.hstack([idx.astype(np.float64), np.ones((idx.shape[0], 1))])
    world = (mask.affine @ homog.T).T[:, :3]
    centroid = world.mean(axis=0)
    if world.shape[0] >= 3:
        cov = np.cov((world - centroid).T)
        evals, evecs = np.linalg.eigh(cov)
        order = np.argsort(evals)[::-1]
        axes = evecs[:, order].T                    # rows, descending variance
    else:
        axes = np.eye(3)
    lengths = []
    for a in axes:
        proj = (world - centroid) @ a
        lengths.append(float(proj.max() - proj.min()))
    voxel_cc = float(np.prod(mask.spacing)) / 1000.0
    return {"centroid": centroid, "axes": axes, "lengths": lengths,
            "volume_cc": round(idx.shape[0] * voxel_cc, 1), "long_axis": axes[0]}


def traversal_length(traj: Trajectory, mask: Volume, n=256) -> float:
    """Length (mm) of the trajectory that lies inside the clot mask."""
    pts = traj.sample(n)
    inside = 0
    for p in pts:
        v = np.round(mask.world_to_voxel(p)).astype(int)
        if mask.in_bounds(v) and mask.data[tuple(v)] > 0:
            inside += 1
    return float(inside) / n * traj.length


@dataclass
class HematomaCorridor:
    entry: np.ndarray
    target: np.ndarray
    traversal_mm: float
    min_clearance_mm: float
    nearest: str
    clearances: dict
    feasible: bool
    band: str
    length_mm: float
    angle_deg: float

    @property
    def trajectory(self) -> Trajectory:
        return Trajectory(self.entry, self.target)


def plan_hematoma(volume: Volume, hematoma_mask: Volume, structures, *,
                  max_angle_deg=60.0, n_theta=9, n_phi=26, n=220,
                  deformation_factory=None):
    """Rank evacuation corridors to the clot centroid by clot traversal (subject
    to safety clearances). Entries are searched in a cone aligned to the clot's
    long axis (toward the nearer surface), so corridors run down the clot.

    Returns (ranked corridors best-first, hematoma_axes dict).
    """
    from scipy.spatial import cKDTree
    ax = hematoma_axes(hematoma_mask)
    if ax["centroid"] is None:
        return [], ax
    C = ax["centroid"]
    a1 = np.asarray(ax["long_axis"], float)
    # orient the long axis toward the nearer surface (away from the volume centre)
    ni, nj, nk = volume.shape
    vol_center = volume.voxel_to_world(((ni - 1) / 2, (nj - 1) / 2, (nk - 1) / 2))
    d = float(np.dot(a1, C - vol_center))
    axis_entry = a1 * (1.0 if d >= 0 else -1.0)
    if abs(d) < 1e-6:
        axis_entry = np.array([0.0, 0.0, 1.0])

    radius = default_entry_radius(volume, C, axis_entry)
    entries = candidate_entries(C, radius, axis=axis_entry, max_angle_deg=max_angle_deg,
                                n_theta=n_theta, n_phi=n_phi)

    trees = None
    if deformation_factory is None:
        trees = {s.name: (cKDTree(_mask_world_points(s.mask)) if s.mask.data.any() else None)
                 for s in structures}

    corridors = []
    for e in np.atleast_2d(entries):
        tr = Trajectory(e, C)
        trav = traversal_length(tr, hematoma_mask, n)
        clears = {}
        for s in structures:
            if deformation_factory is not None:
                pts = _mask_world_points(s.mask, deformation_factory(tr))
                tree = cKDTree(pts) if len(pts) else None
            else:
                tree = trees[s.name]
            clears[s.name] = float(np.min(tree.query(tr.sample(n))[0])) if tree is not None else float("inf")
        nearest = min(clears, key=clears.get) if clears else None
        minc = clears[nearest] if nearest else float("inf")
        feasible = all(clears[s.name] >= s.min_clearance_mm for s in structures) and \
            tr.angle_from_axis(axis_entry) <= max_angle_deg
        corridors.append(HematomaCorridor(
            entry=e, target=C, traversal_mm=round(trav, 1),
            min_clearance_mm=round(minc, 2) if np.isfinite(minc) else float("inf"),
            nearest=nearest, clearances={k: round(v, 2) for k, v in clears.items()},
            feasible=feasible, band=risk_band(minc),
            length_mm=round(tr.length, 1), angle_deg=round(tr.angle_from_axis(), 1)))
    # feasible first, then most clot traversed, then safest
    corridors.sort(key=lambda c: (c.feasible, c.traversal_mm, c.min_clearance_mm), reverse=True)
    return corridors, ax
