"""Deformation-aware structure risk.

Computes trajectory clearance to critical structures both against static anatomy
and against the anatomy as predicted to deform when an instrument is advanced
along the trajectory. The deformation-aware comparison is the key differentiator:
risk is assessed against where tissue is predicted to be, not only where it was
on the preoperative scan.
"""
from __future__ import annotations
import numpy as np
from .planning import risk_band
from . import atlas as A


def select_labels(labels, keywords):
    """Exact label names containing any of the given substrings."""
    return [l for l in labels if any(k.lower() in l.lower() for k in keywords)]


def _world_coords(atlas_vol, labels, names):
    idx = [labels.index(n) for n in names if n in labels]
    AA = np.asarray(atlas_vol.data)
    sel = np.isin(AA, idx)
    coords = np.argwhere(sel)
    if coords.size == 0:
        return np.empty((0, 3)), np.empty((0,), int)
    lab = AA[sel]
    homog = np.hstack([coords.astype(np.float64), np.ones((coords.shape[0], 1))])
    world = (atlas_vol.affine @ homog.T).T[:, :3]
    return world, lab


def _nearest(traj, world, lab, labels, n):
    best, bl = float("inf"), None
    for p in traj.sample(n):
        dd = np.linalg.norm(world - p[None, :], axis=1)
        j = int(np.argmin(dd))
        if dd[j] < best:
            best, bl = float(dd[j]), int(lab[j])
    return best, (labels[bl] if bl is not None else None)


def deformation_aware_risk(traj, atlas_vol, labels, names=None, deformation=None, n=240) -> dict:
    """Static and (optionally) deformation-aware clearance/nearest/band."""
    names = names or [l for l in A.ELOQUENT_DEFAULT if l in labels]
    world, lab = _world_coords(atlas_vol, labels, names)
    if world.shape[0] == 0:
        return {"static_clearance_mm": float("inf"), "static_nearest": None, "static_band": "LOW"}
    sd, sn = _nearest(traj, world, lab, labels, n)
    out = {"static_clearance_mm": round(sd, 2), "static_nearest": sn, "static_band": risk_band(sd)}
    if deformation is not None:
        world_def = world + deformation.field(world)
        dd, dn = _nearest(traj, world_def, lab, labels, n)
        out.update({"deformed_clearance_mm": round(dd, 2), "deformed_nearest": dn,
                    "deformed_band": risk_band(dd), "clearance_change_mm": round(dd - sd, 2)})
    return out
