"""Render the first-order probe-induced tissue-deformation field on a real MRI
slice through a planned trajectory."""
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

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)
ni, nj, nk = vol.shape
target_vox = (ni // 2 + 2, 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)
pd = ProbeDeformation(traj, probe_radius_mm=1.5, influence_radius_mm=18.0)
print("deformation summary:", pd.summary())

# sample the displacement magnitude over the sagittal slice through the target
disp = vol.window().data
xi = int(target_vox[0])
J, K = np.meshgrid(np.arange(nj), np.arange(nk), indexing="ij")
ones = np.ones_like(J)
ijk = np.stack([xi * ones, J, K], axis=-1).reshape(-1, 3).astype(float)
homog = np.hstack([ijk, np.ones((ijk.shape[0], 1))])
world = (vol.affine @ homog.T).T[:, :3]
mag = pd.magnitude(world).reshape(nj, nk)

ev, tv = np.array(entry_vox, float), np.array(target_vox, float)
fig, ax = plt.subplots(1, 2, figsize=(11, 5.4), facecolor="#0b0f18")
for a in ax:
    a.imshow(disp[xi, :, :].T, cmap="gray", origin="lower"); a.axis("off")
ax[0].plot([ev[1], tv[1]], [ev[2], tv[2]], "-", color="#39ff9e", lw=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("Planned trajectory (sagittal)", color="#9fd0ff", fontsize=11)
hm = np.ma.masked_where(mag.T < 0.05, mag.T)
im = ax[1].imshow(hm, cmap="inferno", origin="lower", alpha=0.85)
ax[1].plot([ev[1], tv[1]], [ev[2], tv[2]], "-", color="#39ff9e", lw=1.5)
ax[1].set_title("Predicted tissue displacement (mm)", color="#9fd0ff", fontsize=11)
cb = fig.colorbar(im, ax=ax[1], fraction=0.046, pad=0.04); cb.ax.tick_params(colors="#9fb3d4")
cb.set_label("displacement (mm)", color="#9fb3d4")
fig.suptitle(f"Sulcora — real-time probe-induced deformation (first-order cavity-expansion model)  ·  "
             f"max {pd.max_shift_mm:.1f} mm", color="#dce6f5", fontsize=12)
fig.text(0.5, 0.02, "Analytical real-time approximation for planning preview — not a validated patient-specific FEM/AI model.",
         color="#8aa0c0", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.04, 1, 0.93])
out = os.path.join(DOCS, "deformation_field.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nDEFORMATION PROOF OK")
