"""Finite-element vs. analytical tissue deformation.

Solves the linear-elasticity FEM model of probe-induced tissue displacement on a
real MRI trajectory and compares it against the first-order analytical
cavity-expansion field. Renders (1) the FEM displacement magnitude + vectors on a
slice through the probe, and (2) the radial decay of both models.

RESEARCH SOFTWARE — NOT FOR CLINICAL USE. Demonstration on public MNI152 data.
"""
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.fem import FEMProbeDeformation

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
# a trajectory from a frontal entry toward a deep target
target = vol.voxel_to_world((ni // 2 + 3, nj // 2, nk // 2))
entry = vol.voxel_to_world((int(ni * 0.70), int(nj * 0.62), int(nk * 0.92)))
traj = Trajectory(entry, target)

probe_r, infl = 2.0, 20.0
fem = FEMProbeDeformation(traj, probe_radius_mm=probe_r, influence_radius_mm=infl,
                          poisson=0.45, element_mm=2.5)
ana = ProbeDeformation(traj, probe_radius_mm=probe_r, influence_radius_mm=infl)

print("FEM summary:     ", fem.summary())
print("analytical sumry:", ana.summary())

# ---- radial decay comparison (perpendicular to the probe, at mid-shaft) ---- #
v = traj.target - traj.entry
vhat = v / np.linalg.norm(v)
ref = np.array([1.0, 0, 0]) if abs(vhat[0]) < 0.9 else np.array([0, 1.0, 0])
e1 = np.cross(vhat, ref); e1 /= np.linalg.norm(e1)
mid = traj.entry + 0.5 * v
rr = np.linspace(0.0, infl + 6, 60)
pts = mid[None, :] + rr[:, None] * e1[None, :]
fem_mag = fem.magnitude(pts)
ana_mag = ana.magnitude(pts)
print(f"peak shift  FEM={fem_mag.max():.2f}mm  analytical={ana_mag.max():.2f}mm")

# ---- field figure on a coronal slice through the target --------------------- #
disp = vol.window().data
tv = np.round(vol.world_to_voxel(traj.target)).astype(int)
y = int(np.clip(tv[1], 0, nj - 1))

# sample the FEM displacement on the slice plane (world coords of that y-plane)
xs = np.arange(ni); zs = np.arange(nk)
gx, gz = np.meshgrid(xs, zs, indexing="ij")
plane_vox = np.stack([gx, np.full_like(gx, y), gz], axis=-1).reshape(-1, 3)
ones = np.ones((plane_vox.shape[0], 1))
plane_world = (vol.affine @ np.hstack([plane_vox, ones]).T).T[:, :3]
fem_field = fem.field(plane_world)
fem_slice = np.linalg.norm(fem_field, axis=1).reshape(ni, nk)

fig = plt.figure(figsize=(12.5, 5.2), facecolor="#0b0f18")
ax0 = fig.add_subplot(1, 2, 1)
ax0.imshow(disp[:, y, :].T, cmap="gray", origin="lower")
m = np.ma.masked_where(fem_slice.T < 0.05, fem_slice.T)
im = ax0.imshow(m, cmap="inferno", alpha=0.7, origin="lower")
ev, tw = vol.world_to_voxel(traj.entry), vol.world_to_voxel(traj.target)
ax0.plot([ev[0], tw[0]], [ev[2], tw[2]], "-", color="#39ff9e", lw=2.0)
ax0.set_title("FEM tissue displacement (coronal)", color="#9fd0ff", fontsize=11)
ax0.axis("off")
cb = fig.colorbar(im, ax=ax0, fraction=0.046, pad=0.04); cb.set_label("|u| (mm)", color="#cfe2ff")
cb.ax.yaxis.set_tick_params(color="#cfe2ff"); plt.setp(plt.getp(cb.ax, "yticklabels"), color="#cfe2ff")

ax1 = fig.add_subplot(1, 2, 2, facecolor="#0e1626")
ax1.plot(rr, fem_mag, "-", color="#ff7f3f", lw=2.2, label="FEM (linear elasticity)")
ax1.plot(rr, ana_mag, "--", color="#6fb6ff", lw=2.0, label="analytical (cavity expansion)")
ax1.axvline(probe_r, color="#888", ls=":", lw=1)
ax1.set_xlabel("distance from probe axis (mm)", color="#cfe2ff")
ax1.set_ylabel("displacement |u| (mm)", color="#cfe2ff")
ax1.set_title("Radial decay: FEM vs. analytical", color="#9fd0ff", fontsize=11)
ax1.tick_params(colors="#9fb3d0"); ax1.grid(alpha=0.15)
for s in ax1.spines.values():
    s.set_color("#33425f")
ax1.legend(facecolor="#0e1626", edgecolor="#33425f", labelcolor="#cfe2ff")

fig.suptitle("Finite-element tissue-deformation model (Navier-Cauchy, hex FEM)", color="#dce6f5", fontsize=13)
fig.text(0.5, 0.015, "Solve-domain far field held fixed (skull); probe wall pushed outward. "
         "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.05, 1, 0.94])
out = os.path.join(DOCS, "fem_deformation.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nFEM PROOF OK")
