"""Demonstrate deformation-aware risk: static clearance vs. clearance against
predicted deformed anatomy, on a real MRI with cortical + subcortical structures."""
import os, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
from sulcora.io import load
from sulcora.planning import Trajectory
from sulcora.deformation import ProbeDeformation
from sulcora import atlas as A
from sulcora import risk as R

DOCS = os.path.join(HERE, "docs", "images"); os.makedirs(DOCS, exist_ok=True)
nii = os.path.join(HERE, "sample_data", "mni152_t1_2mm.nii.gz")
if not os.path.exists(nii):
    import nibabel as nib
    from nilearn.datasets import load_mni152_template
    nib.save(load_mni152_template(resolution=2), nii)

vol = load(nii)
cort, clabels = A.load_harvard_oxford("cort-maxprob-thr25-2mm")
cort = A.resample_atlas_to(cort, vol)
sub, slabels = A.load_harvard_oxford("sub-maxprob-thr25-2mm")
sub = A.resample_atlas_to(sub, vol)
elo = [l for l in A.ELOQUENT_DEFAULT if l in clabels]
deep = R.select_labels(slabels, A.SUBCORTICAL_DEFAULT)
print("cortical eloquent:", len(elo), "| subcortical:", len(deep))

ni, nj, nk = vol.shape
target_vox = (ni // 2 + 4, nj // 2, nk // 2)
entry_vox = (int(ni * 0.70), int(nj * 0.60), int(nk * 0.92))
tw = vol.voxel_to_world(target_vox); ew = vol.voxel_to_world(entry_vox)
traj = Trajectory(ew, tw)
pd = ProbeDeformation(traj, probe_radius_mm=2.0, influence_radius_mm=20.0)

rep_c = R.deformation_aware_risk(traj, cort, clabels, names=elo, deformation=pd)
rep_s = R.deformation_aware_risk(traj, sub, slabels, names=deep, deformation=pd)
print("CORTICAL:", rep_c)
print("SUBCORTICAL:", rep_s)

# figure: numbers panel + eloquent overlay
disp = vol.window().data
emask = A.structure_mask(cort, clabels, elo).data
smask = A.structure_mask(sub, slabels, deep).data
tv, ev = np.array(target_vox, float), np.array(entry_vox, float)
xi = int(target_vox[0])

fig, ax = plt.subplots(1, 2, figsize=(12, 5.6), facecolor="#0b0f18",
                       gridspec_kw={"width_ratios": [1.25, 1]})
a0 = ax[0]
a0.imshow(disp[xi, :, :].T, cmap="gray", origin="lower")
for m, c in [(emask, "cool"), (smask, "autumn")]:
    mm = np.ma.masked_where(m[xi, :, :].T == 0, m[xi, :, :].T)
    a0.imshow(mm, cmap=c, alpha=0.45, origin="lower")
a0.plot([ev[1], tv[1]], [ev[2], tv[2]], "-", color="#39ff9e", lw=2.2)
a0.plot(tv[1], tv[2], "o", color="#ff5d5d"); a0.plot(ev[1], ev[2], "o", color="#39ff9e")
a0.set_title("Trajectory · eloquent (cyan) + subcortical (orange)", color="#9fd0ff", fontsize=11); a0.axis("off")

a1 = ax[1]; a1.axis("off"); a1.set_facecolor("#0b0f18")
def band_col(b): return {"HIGH": "#ff5d5d", "CAUTION": "#f5b041", "LOW": "#2ecc71"}[b]
lines = [
    ("Structure-aware risk", "", "#dce6f5", 15, True),
    ("", "", "", 6, False),
    (f"Nearest (cortical): {rep_c['static_nearest']}", "", "#9fb3d4", 12, False),
    (f"  static clearance", f"{rep_c['static_clearance_mm']:.1f} mm  ({rep_c['static_band']})", band_col(rep_c['static_band']), 13, False),
    (f"  deformation-aware", f"{rep_c['deformed_clearance_mm']:.1f} mm  ({rep_c['deformed_band']})", band_col(rep_c['deformed_band']), 13, False),
    (f"  change under predicted shift", f"{rep_c['clearance_change_mm']:+.1f} mm", "#9fb3d4", 12, False),
    ("", "", "", 8, False),
    (f"Nearest (subcortical): {rep_s['static_nearest']}", "", "#9fb3d4", 12, False),
    (f"  static clearance", f"{rep_s['static_clearance_mm']:.1f} mm  ({rep_s['static_band']})", band_col(rep_s['static_band']), 13, False),
    (f"  deformation-aware", f"{rep_s['deformed_clearance_mm']:.1f} mm  ({rep_s['deformed_band']})", band_col(rep_s['deformed_band']), 13, False),
]
y = 0.96
for left, right, col, sz, bold in lines:
    if left or right:
        a1.text(0.02, y, left, color=col, fontsize=sz, fontweight=("bold" if bold else "normal"), transform=a1.transAxes, va="top")
        if right:
            a1.text(0.98, y, right, color=col, fontsize=sz, fontweight="bold", transform=a1.transAxes, va="top", ha="right")
    y -= 0.072 + (sz / 900.0)
fig.suptitle("Sulcora v0.4 — deformation-aware risk (static vs. predicted deformed anatomy)",
             color="#dce6f5", fontsize=13)
fig.text(0.5, 0.02, "Risk recomputed against where tissue is predicted to be after probe-induced shift — not only the preoperative scan.",
         color="#8aa0c0", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.04, 1, 0.93])
out = os.path.join(DOCS, "deformation_aware_risk.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nDEF-RISK PROOF OK")
