"""Demonstrate real, atlas-based structure-aware risk on a planned trajectory,
and render an overlay figure (eloquent cortex + trajectory + risk)."""
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 import atlas as A

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
    os.makedirs(os.path.dirname(nii), exist_ok=True)
    nib.save(load_mni152_template(resolution=2), nii)

vol = load(nii)
atlas_vol, labels = A.load_harvard_oxford()
atlas_vol = A.resample_atlas_to(atlas_vol, vol)          # onto the patient grid
elo = [l for l in A.ELOQUENT_DEFAULT if l in labels]
print("eloquent structures:", len(elo))

# plan a trajectory toward a deep target, entry near motor strip
ni, nj, nk = vol.shape
target_vox = (ni // 2 + 3, nj // 2, nk // 2)
entry_vox = (int(ni * 0.66), int(nj * 0.62), int(nk * 0.92))
tw = vol.voxel_to_world(target_vox); ew = vol.voxel_to_world(entry_vox)
traj = Trajectory(ew, tw)
rep = A.structure_risk(traj, atlas_vol, labels, elo)
print(f"length={traj.length:.1f}mm angle={traj.angle_from_axis():.0f}deg")
print(f"RISK: nearest={rep['nearest']}  clearance={rep['clearance_mm']:.1f}mm  band={rep['band']}")

# overlay figure
disp = vol.window().data
emask = A.structure_mask(atlas_vol, labels, elo).data
tv, ev = np.array(target_vox, float), np.array(entry_vox, float)
band_col = {"HIGH": "#ff5d5d", "CAUTION": "#f5b041", "LOW": "#2ecc71"}[rep["band"]]

def overlay(ax, bg, mask, line, pts, title):
    ax.imshow(bg.T, cmap="gray", origin="lower")
    m = np.ma.masked_where(mask.T == 0, mask.T)
    ax.imshow(m, cmap="cool", alpha=0.45, origin="lower")
    ax.plot(*line, "-", color=band_col, lw=2.4)
    ax.plot(*pts[0], "o", color="#39ff9e"); ax.plot(*pts[1], "o", color="#ff5d5d")
    ax.set_title(title, color="#9fd0ff", fontsize=11); ax.axis("off")

fig, ax = plt.subplots(1, 3, figsize=(13, 4.7), facecolor="#0b0f18")
z, y, x = int(tv[2]), int(tv[1]), int(tv[0])
overlay(ax[0], disp[:, :, z], emask[:, :, z], ([ev[0], tv[0]], [ev[1], tv[1]]),
        ([ev[0], ev[1]], [tv[0], tv[1]]), "Axial")
overlay(ax[1], disp[:, y, :], emask[:, y, :], ([ev[0], tv[0]], [ev[2], tv[2]]),
        ([ev[0], ev[2]], [tv[0], tv[2]]), "Coronal")
overlay(ax[2], disp[x, :, :], emask[x, :, :], ([ev[1], tv[1]], [ev[2], tv[2]]),
        ([ev[1], ev[2]], [tv[1], tv[2]]), "Sagittal")
fig.suptitle(f"Structure-aware risk (real Harvard-Oxford atlas) — nearest: {rep['nearest']}  ·  "
             f"clearance {rep['clearance_mm']:.1f} mm  ·  {rep['band']}",
             color="#dce6f5", fontsize=12)
fig.text(0.5, 0.015, "Magenta = eloquent cortex (motor/speech/language/visual). Trajectory coloured by risk band.",
         color="#8aa0c0", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.04, 1, 0.93])
out = os.path.join(DOCS, "risk_overlay.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nRISK PROOF OK")
