"""White-matter tract avoidance: overlay real fiber tracts (Juelich atlas) and a
trajectory, reporting clearance to the corticospinal tract."""
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
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)
jul, jlab = A.load_juelich()
jul = A.resample_atlas_to(jul, vol)
tracts = R.select_labels(jlab, ["Corticospinal", "Optic radiation"])
print("tracts used:", tracts)

ni, nj, nk = vol.shape
tw = vol.voxel_to_world((ni // 2 + 6, nj // 2 + 4, nk // 2))
ew = vol.voxel_to_world((int(ni * 0.72), int(nj * 0.6), int(nk * 0.9)))
traj = Trajectory(ew, tw)
rep = R.deformation_aware_risk(traj, jul, jlab, names=tracts)
print("clearance to tracts:", rep["static_clearance_mm"], "mm — nearest:", rep["static_nearest"])

disc = vol.window().data
cst = A.structure_mask(jul, jlab, R.select_labels(jlab, ["Corticospinal"])).data
opt = A.structure_mask(jul, jlab, R.select_labels(jlab, ["Optic radiation"])).data
tv, ev = np.array(vol.world_to_voxel(tw)), np.array(vol.world_to_voxel(ew))
xi = int(round(tv[0]))
fig, ax = plt.subplots(1, 2, figsize=(11, 5.6), facecolor="#0b0f18")
for a, (m1lab) in zip(ax, ["sag", "cor"]):
    a.axis("off")
img = disc[xi, :, :].T
ax[0].imshow(img, cmap="gray", origin="lower")
for m, c in [(cst[xi, :, :].T, "autumn"), (opt[xi, :, :].T, "winter")]:
    ax[0].imshow(np.ma.masked_where(m == 0, m), cmap=c, alpha=0.6, origin="lower")
ax[0].plot([ev[1], tv[1]], [ev[2], tv[2]], "-", color="#39ff9e", lw=2.2)
ax[0].plot(tv[1], tv[2], "o", color="#ff5d5d"); ax[0].plot(ev[1], ev[2], "o", color="#39ff9e")
ax[0].set_title("Sagittal — corticospinal tract (red) + optic radiation (blue)", color="#9fd0ff", fontsize=10.5)
yj = int(round(tv[1]))
ax[1].imshow(disc[:, yj, :].T, cmap="gray", origin="lower")
for m, c in [(cst[:, yj, :].T, "autumn"), (opt[:, yj, :].T, "winter")]:
    ax[1].imshow(np.ma.masked_where(m == 0, m), cmap=c, alpha=0.6, origin="lower")
ax[1].plot([ev[0], tv[0]], [ev[2], tv[2]], "-", color="#39ff9e", lw=2.2)
ax[1].plot(tv[0], tv[2], "o", color="#ff5d5d"); ax[1].plot(ev[0], ev[2], "o", color="#39ff9e")
ax[1].set_title("Coronal", color="#9fd0ff", fontsize=10.5)
fig.suptitle(f"Sulcora v0.7 — white-matter tract avoidance (Juelich atlas) · "
             f"clearance to {rep['static_nearest']}: {rep['static_clearance_mm']:.1f} mm",
             color="#dce6f5", fontsize=12)
fig.text(0.5, 0.02, "Real fiber tracts now scored alongside cortical and deep structures.",
         color="#8aa0c0", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.04, 1, 0.93])
out = os.path.join(DOCS, "tract_avoidance.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nTRACT PROOF OK")
