"""Fetch a real MRI (MNI152), exercise the Sulcora core end-to-end, and render
proof figures (tri-planar MPR + 3D surface with a planned trajectory)."""
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SAMPLE = os.path.join(HERE, "sample_data")
DOCS = os.path.join(HERE, "docs", "images")
os.makedirs(SAMPLE, exist_ok=True)
os.makedirs(DOCS, exist_ok=True)

# ---- 1. obtain a real MRI and save as NIfTI (proves the package's sample data) ----
nii_path = os.path.join(SAMPLE, "mni152_t1_2mm.nii.gz")
if not os.path.exists(nii_path):
    import nibabel as nib
    from nilearn.datasets import load_mni152_template
    img = load_mni152_template(resolution=2)
    nib.save(img, nii_path)
print("sample MRI:", nii_path)

# ---- 2. load through the real package IO ----
import sys
sys.path.insert(0, HERE)
from sulcora.io import load
from sulcora.volume import Volume
from sulcora.planning import Trajectory, risk_band

vol = load(nii_path)
print("loaded volume shape", vol.shape, "spacing", np.round(vol.spacing, 2))

# ---- 3. define a plan in world (mm) coords; compute metrics ----
ni, nj, nk = vol.shape
target_vox = (ni // 2 + 2, nj // 2, nk // 2)            # ~central deep target
entry_vox = (ni // 2 + 2, int(nj * 0.80), int(nk * 0.92))  # frontal entry
target_w = vol.voxel_to_world(target_vox)
entry_w = vol.voxel_to_world(entry_vox)
traj = Trajectory(entry_w, target_w)

# example critical ROI mask (sphere) to demonstrate the clearance metric
roi = Volume(np.zeros(vol.shape, np.uint8), vol.affine.copy())
cx, cy, cz = ni // 2 + 6, nj // 2 + 4, nk // 2
ii, jj, kk = np.mgrid[0:ni, 0:nj, 0:nk]
roi.data[((ii - cx) ** 2 + (jj - cy) ** 2 + (kk - cz) ** 2) <= 9] = 1
clear_mm, _ = traj.min_clearance_to_mask(roi, n=180)

print(f"plan: length={traj.length:.1f} mm  angle={traj.angle_from_axis():.0f} deg  "
      f"clearance(example ROI)={clear_mm:.1f} mm  band={risk_band(clear_mm)}")

# ---- 4. tri-planar MPR figure with trajectory ----
disp = vol.window().data
tv, ev = np.array(target_vox, float), np.array(entry_vox, float)

fig, ax = plt.subplots(1, 3, figsize=(13, 4.6), facecolor="#0b0f18")
def style(a, title):
    a.set_title(title, color="#9fd0ff", fontsize=11); a.axis("off")
# axial (k fixed)
ax[0].imshow(disp[:, :, int(tv[2])].T, cmap="gray", origin="lower")
ax[0].plot([ev[0], tv[0]], [ev[1], tv[1]], "-", color="#39ff9e", lw=2)
ax[0].plot(tv[0], tv[1], "o", color="#ff5d5d"); ax[0].plot(ev[0], ev[1], "o", color="#39ff9e")
style(ax[0], "Axial")
# coronal (j fixed)
ax[1].imshow(disp[:, int(tv[1]), :].T, cmap="gray", origin="lower")
ax[1].plot([ev[0], tv[0]], [ev[2], tv[2]], "-", color="#39ff9e", lw=2)
ax[1].plot(tv[0], tv[2], "o", color="#ff5d5d"); ax[1].plot(ev[0], ev[2], "o", color="#39ff9e")
style(ax[1], "Coronal")
# sagittal (i fixed)
ax[2].imshow(disp[int(tv[0]), :, :].T, cmap="gray", origin="lower")
ax[2].plot([ev[1], tv[1]], [ev[2], tv[2]], "-", color="#39ff9e", lw=2)
ax[2].plot(tv[1], tv[2], "o", color="#ff5d5d"); ax[2].plot(ev[1], ev[2], "o", color="#39ff9e")
style(ax[2], "Sagittal")
fig.suptitle(f"Sulcora Planner — trajectory {traj.length:.0f} mm, {traj.angle_from_axis():.0f}° from vertical "
             f"(real MNI152 MRI)", color="#dce6f5", fontsize=12)
fig.tight_layout(rect=[0, 0, 1, 0.94])
mpr_png = os.path.join(DOCS, "mpr_trajectory.png")
fig.savefig(mpr_png, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", mpr_png)

# ---- 5. 3D surface figure with trajectory ----
from skimage.measure import marching_cubes
from scipy.ndimage import gaussian_filter
sm = gaussian_filter(disp.astype(float), 1.0)
verts, faces, _, _ = marching_cubes(sm, level=40, spacing=vol.spacing)
fig2 = plt.figure(figsize=(7.5, 7.0), facecolor="#0b0f18")
a3 = fig2.add_subplot(111, projection="3d"); a3.set_facecolor("#0b0f18")
mesh = Poly3DCollection(verts[faces][::3], alpha=0.18)
mesh.set_facecolor((0.80, 0.71, 0.68)); mesh.set_edgecolor("none")
a3.add_collection3d(mesh)
# trajectory in the same spacing-scaled coords
ew = ev * vol.spacing; twv = tv * vol.spacing
a3.plot([ew[0], twv[0]], [ew[1], twv[1]], [ew[2], twv[2]], color="#39ff9e", lw=3)
a3.scatter(*twv, color="#ff5d5d", s=40); a3.scatter(*ew, color="#39ff9e", s=40)
mx = verts.max(0)
a3.set_xlim(0, mx[0]); a3.set_ylim(0, mx[1]); a3.set_zlim(0, mx[2])
a3.set_box_aspect((mx[0], mx[1], mx[2])); a3.axis("off")
a3.view_init(elev=12, azim=-70)
a3.set_title("3D reconstruction with planned trajectory", color="#9fd0ff")
surf_png = os.path.join(DOCS, "surface_trajectory.png")
fig2.savefig(surf_png, dpi=130, facecolor=fig2.get_facecolor()); plt.close(fig2)
print("wrote", surf_png)
print("PROOF OK")
