"""Robotics-validation demo: simulate a robot executing a planned trajectory on
the virtual brain and score targeting accuracy + deformation-aware safety."""
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.robot import RoboticExecution
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)
cort, clabels = A.load_harvard_oxford("cort-maxprob-thr25-2mm")
cort = A.resample_atlas_to(cort, vol)
elo = [l for l in A.ELOQUENT_DEFAULT if l in clabels]

# a planned trajectory (deep target, frontal entry)
ni, nj, nk = vol.shape
tw = vol.voxel_to_world((ni // 2 + 3, nj // 2, nk // 2))
ew = vol.voxel_to_world((int(ni * 0.66), int(nj * 0.60), int(nk * 0.92)))
planned = Trajectory(ew, tw)

# deformation-aware clearance function for the safety score (expensive)
def clearance_fn(traj):
    pd = ProbeDeformation(traj, probe_radius_mm=2.0, influence_radius_mm=20.0)
    rep = R.deformation_aware_risk(traj, cort, clabels, names=elo, deformation=pd, n=70)
    return rep.get("deformed_clearance_mm", rep["static_clearance_mm"])

robot = RoboticExecution(planned, targeting_sigma_mm=1.0, angular_sigma_deg=1.0)
s = robot.simulate(n=1500, tre_tol_mm=2.0, seed=7)                       # accuracy (fast)
ss = robot.simulate(n=120, clearance_fn=clearance_fn, tre_tol_mm=2.0,    # safety (subset)
                    clearance_min_mm=3.0, seed=7)
s.update({k: ss[k] for k in ("mean_clearance_mm", "min_clearance_mm", "clearance_min_mm", "pass_pct")})
print({k: v for k, v in s.items() if not k.startswith("_")})

# figure: scatter of achieved targeting offsets + metrics
off = s["_offsets"]; tre = s["_tre"]
fig, ax = plt.subplots(1, 2, figsize=(12, 5.6), facecolor="#0b0f18",
                       gridspec_kw={"width_ratios": [1, 1]})
a0 = ax[0]; a0.set_facecolor("#0b0f18")
within = tre <= s["tre_tol_mm"]
a0.scatter(off[~within, 0], off[~within, 1], s=8, c="#ff5d5d", alpha=.6, label="outside tol")
a0.scatter(off[within, 0], off[within, 1], s=8, c="#39ff9e", alpha=.5, label="within tol")
th = np.linspace(0, 2*np.pi, 100)
a0.plot(s["tre_tol_mm"]*np.cos(th), s["tre_tol_mm"]*np.sin(th), "--", color="#9fd0ff", lw=1.5)
a0.set_aspect("equal"); a0.set_xlabel("lateral error (mm)", color="#9fb3d4"); a0.set_ylabel("lateral error (mm)", color="#9fb3d4")
a0.set_title("Achieved targeting over 1,500 robotic runs", color="#9fd0ff", fontsize=11)
a0.tick_params(colors="#6f84a6"); a0.legend(facecolor="#101a2c", edgecolor="#2a3a5c", labelcolor="#dce6f5", fontsize=9)
for sp in a0.spines.values(): sp.set_color("#2a3a5c")

a1 = ax[1]; a1.axis("off")
rows = [
    ("Planned depth", f"{planned.length:.0f} mm"),
    ("Entry placement noise (1σ)", f"{robot.targeting_sigma_mm:.1f} mm"),
    ("Angular noise (1σ)", f"{robot.angular_sigma_deg:.1f}°"),
    ("", ""),
    ("Mean targeting error (TRE)", f"{s['mean_tre_mm']:.2f} mm"),
    ("95th-percentile TRE", f"{s['p95_tre_mm']:.2f} mm"),
    ("Worst-case TRE", f"{s['max_tre_mm']:.2f} mm"),
    (f"Within {s['tre_tol_mm']:.0f} mm tolerance", f"{s['tre_within_tol_pct']:.1f}%"),
    ("", ""),
    ("Mean clearance (deformation-aware)", f"{s['mean_clearance_mm']:.1f} mm"),
    ("Worst clearance", f"{s['min_clearance_mm']:.1f} mm"),
    (f"PASS (accuracy & safety)", f"{s['pass_pct']:.1f}%"),
]
y = 0.95
for k, v in rows:
    if k:
        hot = k.startswith("PASS")
        a1.text(0.02, y, k, color=("#dce6f5" if hot else "#9fb3d4"), fontsize=(13 if hot else 12),
                fontweight=("bold" if hot else "normal"), transform=a1.transAxes, va="top")
        a1.text(0.98, y, v, color=("#2ecc71" if hot else "#dce6f5"), fontsize=(14 if hot else 12),
                fontweight="bold", transform=a1.transAxes, va="top", ha="right")
    y -= 0.078
fig.suptitle("Sulcora v0.5 — robotics-validation environment (simulated robotic execution)",
             color="#dce6f5", fontsize=13)
fig.text(0.5, 0.02, "Virtual-brain test bench: score a robot's targeting accuracy and deformation-aware safety before any patient contact.",
         color="#8aa0c0", ha="center", fontsize=9)
fig.tight_layout(rect=[0, 0.04, 1, 0.93])
out = os.path.join(DOCS, "robotics_validation.png")
fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
print("wrote", out, "\nROBOT PROOF OK")
