"""GUI-agnostic planning-session helpers.

Thin, unit-tested glue the desktop app uses to turn a target + a set of critical
structures into a live risk readout and an optimizer recommendation. Kept out of
``app.py`` (which needs a display) so the logic is testable headlessly.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band
from .vessels import _mask_world_points
from .optimizer import plan_safest


def clearance_field(traj: Trajectory, mask: Volume, n=256) -> Volume:
    """Per-voxel distance (mm) from each structure voxel to the trajectory.

    Returns a Volume the shape of `mask` holding, at each set voxel, its minimum
    distance to the trajectory (0 elsewhere). Intended for a proximity heatmap —
    e.g. colour the structure grey (far / safe) → red (close / dangerous).
    """
    from scipy.spatial import cKDTree
    idx = np.argwhere(np.asarray(mask.data) > 0)
    field = np.zeros(mask.shape, np.float32)
    if idx.size == 0:
        return Volume(field, mask.affine.copy())
    homog = np.hstack([idx.astype(np.float64), np.ones((idx.shape[0], 1))])
    world = (mask.affine @ homog.T).T[:, :3]
    d, _ = cKDTree(traj.sample(n)).query(world)
    field[tuple(idx.T)] = d.astype(np.float32)
    return Volume(field, mask.affine.copy())


def default_entry_radius(volume, target, axis=(0, 0, 1), fraction=0.9, minimum=20.0) -> float:
    """Heuristic radius (mm) for candidate entries: distance from `target` to the
    volume's far surface along `axis`, so entries land near the scalp."""
    a = np.asarray(axis, float); a = a / np.linalg.norm(a)
    ni, nj, nk = volume.shape
    corners = np.array([volume.voxel_to_world((i, j, k))
                        for i in (0, ni - 1) for j in (0, nj - 1) for k in (0, nk - 1)])
    along = corners @ a
    reach = float(np.max(along) - np.dot(np.asarray(target, float), a))
    return max(minimum, fraction * abs(reach))


def risk_summary(traj: Trajectory, structures, deformation=None, n=200) -> dict:
    """Per-structure clearance + overall nearest/band for a trajectory.

    When `deformation` is given, also reports the clearance against the predicted
    deformed anatomy and the change vs. static.
    """
    from scipy.spatial import cKDTree
    per, mins = {}, {}
    for s in structures:
        pts = _mask_world_points(s.mask)
        if len(pts):
            cl = float(np.min(cKDTree(pts).query(traj.sample(n))[0]))
        else:
            cl = float("inf")
        mins[s.name] = cl
        rec = {"clearance_mm": round(cl, 2) if np.isfinite(cl) else float("inf"),
               "band": risk_band(cl), "min_clearance_mm": s.min_clearance_mm}
        if deformation is not None and len(pts):
            dcl = float(np.min(cKDTree(_mask_world_points(s.mask, deformation)).query(traj.sample(n))[0]))
            rec["deformed_clearance_mm"] = round(dcl, 2)
            rec["deformed_band"] = risk_band(dcl)
            rec["clearance_change_mm"] = round(dcl - cl, 2)
        per[s.name] = rec

    out = {"per_structure": per}
    if mins:
        nearest = min(mins, key=mins.get)
        out.update({"nearest": nearest,
                    "min_clearance_mm": round(mins[nearest], 2) if np.isfinite(mins[nearest]) else float("inf"),
                    "band": risk_band(mins[nearest])})
        if deformation is not None:
            dmins = {k: v.get("deformed_clearance_mm") for k, v in per.items()
                     if v.get("deformed_clearance_mm") is not None}
            if dmins:
                dn = min(dmins, key=dmins.get)
                out["deformed_min_clearance_mm"] = dmins[dn]
                out["deformed_band"] = risk_band(dmins[dn])
    else:
        out.update({"nearest": None, "min_clearance_mm": float("inf"), "band": "LOW"})
    return out


def clearance_profile(traj: Trajectory, structures, n=80):
    """Clearance (mm) to the nearest critical structure at each depth along the
    trajectory. Returns (depths_mm, min_clearance, per_structure) where depths
    run 0 (entry) → length (target); the profile is what a surgeon reads as
    'how close do I get, and where'."""
    from scipy.spatial import cKDTree
    pts = traj.sample(n)
    depths = np.linspace(0.0, traj.length, n)
    best = np.full(n, np.inf)
    per = {}
    for s in structures:
        w = _mask_world_points(s.mask)
        if len(w) == 0:
            per[s.name] = np.full(n, np.inf)
            continue
        d, _ = cKDTree(w).query(pts)
        per[s.name] = d
        best = np.minimum(best, d)
    return depths, best, per


def recommend(volume, target, structures, axis=(0, 0, 1), max_angle_deg=60.0,
              n_theta=8, n_phi=24, max_depth_mm=None, deformation_factory=None, n=200):
    """Run the optimizer over an auto-sized cap of candidate entries around
    `target`; returns ranked TrajectoryPlan list (best first)."""
    radius = default_entry_radius(volume, target, axis)
    return plan_safest(target, structures, radius=radius, axis=axis,
                       max_angle_deg=max_angle_deg, n_theta=n_theta, n_phi=n_phi,
                       max_depth_mm=max_depth_mm, deformation_factory=deformation_factory, n=n)


# Clinically established / FDA-recognised DBS targets only (current literature) —
# not the whole atlas, and not research-only nuclei.
DBS_TARGETS = [
    {"name": "STN (subthalamic nucleus)", "indication": "Parkinson's disease"},
    {"name": "GPi (globus pallidus internus)", "indication": "Parkinson's disease, dystonia"},
    {"name": "VIM (ventral intermediate n. of thalamus)", "indication": "essential / Parkinsonian tremor"},
    {"name": "ANT (anterior nucleus of thalamus)", "indication": "drug-resistant epilepsy (FDA-approved, SANTE)"},
    {"name": "VC/VS (ventral capsule / ventral striatum)", "indication": "OCD (FDA humanitarian device exemption)"},
]


def procedure_targets(volume, procedure, lesion_mask=None) -> dict:
    """Return the target list appropriate to a procedure — the key point that
    targets are procedure-specific, not one global list:

    * ICH  -> detected haematoma(s) (segment CT), each with volume + location;
    * DBS  -> the standard functional nuclei (STN, GPi, VIM, ...);
    * EVD  -> the ventricle frontal horns (right/left);
    * biopsy/LITT/lesion/GBM -> the supplied segmented lesion's centroid;
    * SEEG -> any named structure/target.

    Returns {"kind": ..., "items": [...]}.
    """
    p = str(procedure).lower()
    if p == "ich":
        from .hematoma import list_hematomas
        return {"kind": "detected", "items": list_hematomas(volume)}
    if p == "dbs":
        return {"kind": "named", "items": list(DBS_TARGETS)}
    if p == "evd":
        from .ventricle import segment_ventricles, frontal_horn_target
        vent = segment_ventricles(volume)
        items = []
        for s in ("right", "left"):
            c = frontal_horn_target(vent, s)
            if c is not None:
                items.append({"name": f"Frontal horn ({s})", "centroid": c})
        return {"kind": "detected", "items": items}
    if p in ("lesion", "gbm", "biopsy", "litt", "tumor"):
        if lesion_mask is not None:
            return {"kind": "detected", "items": [lesion_info(lesion_mask)]}
        return {"kind": "needs_mask", "items": []}
    if p == "seeg":
        return {"kind": "named", "items": [{"name": "(any structure / custom target)"}]}
    return {"kind": "unknown", "items": []}


def plan_multi(volume, targets, structures, min_separation_mm=6.0, **kw):
    """Plan multiple non-conflicting corridors (e.g. SEEG electrodes): for each
    target, pick the best feasible corridor that stays at least `min_separation_mm`
    from every already-accepted trajectory. Returns a list aligned with `targets`
    (None if no corridor could be placed)."""
    from scipy.spatial import cKDTree

    def separation(a, b, n=120):
        return float(np.min(cKDTree(b.sample(n)).query(a.sample(n))[0]))

    accepted = []
    for tgt in np.atleast_2d(np.asarray(targets, float)):
        ranked = recommend(volume, tgt, structures, **kw)
        chosen = None
        for p in ranked:
            if not p.feasible:
                continue
            if all(separation(p.trajectory, a.trajectory) >= min_separation_mm for a in accepted):
                chosen = p
                break
        if chosen is None and ranked:
            chosen = ranked[0]
        if chosen is not None:
            accepted.append(chosen)
    return accepted


def lesion_info(mask) -> dict:
    """Centroid (world mm), voxel count, and volume (cc) of a lesion mask —
    e.g. an intracerebral haematoma, abscess, or tumour to be reached/evacuated."""
    idx = np.argwhere(np.asarray(mask.data) > 0)
    if idx.size == 0:
        return {"centroid": None, "voxels": 0, "volume_cc": 0.0}
    homog = np.hstack([idx.astype(np.float64), np.ones((idx.shape[0], 1))])
    world = (mask.affine @ homog.T).T[:, :3]
    voxel_cc = float(np.prod(mask.spacing)) / 1000.0        # mm^3 -> cc
    return {"centroid": world.mean(axis=0), "voxels": int(idx.shape[0]),
            "volume_cc": round(idx.shape[0] * voxel_cc, 1)}


def plan_to_lesion(volume, lesion_mask, structures, **kw):
    """Best corridor to a lesion (targets its centroid). Generalises the planner
    beyond DBS/SEEG to lesion procedures such as haematoma evacuation, biopsy, or
    catheter/laser placement. Returns (ranked plans, lesion_info)."""
    info = lesion_info(lesion_mask)
    if info["centroid"] is None:
        return [], info
    return recommend(volume, info["centroid"], structures, **kw), info
