"""Colour a critical structure grey->red by its clearance to the trajectory.

Instead of a single clearance number, each voxel of a critical structure is
shaded by how close the trajectory comes to it — grey where it is safely distant,
red where the trajectory grazes it — so the surgeon sees *where* the risk is.

RESEARCH SOFTWARE — NOT FOR CLINICAL USE.
"""
import os, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora import session


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    shape = (90, 90, 90)
    ni, nj, nk = shape
    ii, jj, kk = np.meshgrid(np.arange(ni), np.arange(nj), np.arange(nk), indexing="ij")
    c = np.array(shape) / 2
    ell = ((ii - c[0]) / (0.46 * ni)) ** 2 + ((jj - c[1]) / (0.44 * nj)) ** 2 + ((kk - c[2]) / (0.46 * nk)) ** 2
    brain = (40.0 / (1.0 + np.exp((ell - 1.0) / 0.04))).astype(np.float32)
    vol = Volume(brain, np.eye(4))

    # a spherical "eloquent" structure the trajectory grazes on one side
    cen = np.array([55.0, 45.0, 50.0]); R = 12.0
    struct = (((ii - cen[0]) ** 2 + (jj - cen[1]) ** 2 + (kk - cen[2]) ** 2) <= R ** 2).astype(np.uint8)
    mask = Volume(struct, np.eye(4))

    traj = Trajectory(entry=np.array([40.0, 45.0, 82.0]), target=np.array([48.0, 45.0, 40.0]))
    field = session.clearance_field(traj, mask, n=300)
    vals = field.data[struct > 0]
    print(f"structure voxels={int(struct.sum())}  clearance min={vals.min():.1f}mm max={vals.max():.1f}mm")

    THRESH = 12.0                              # >= THRESH mm -> fully grey (safe)
    cmap = LinearSegmentedColormap.from_list("clr", ["#ff2b2b", "#c76b52", "#8f8f8f"])

    disp = vol.window().data
    y = int(round(cen[1]))                     # coronal slice through the structure
    bg = disp[:, y, :]
    fld = field.data[:, y, :]
    ev, tv = traj.entry, traj.target

    fig, ax = plt.subplots(figsize=(6.6, 6.2), facecolor="#0b0f18")
    ax.imshow(bg.T, cmap="gray", origin="lower")
    masked = np.ma.masked_where(struct[:, y, :].T == 0, fld.T)
    im = ax.imshow(masked, cmap=cmap, vmin=0, vmax=THRESH, origin="lower", alpha=0.95)
    ax.plot([ev[0], tv[0]], [ev[2], tv[2]], "-", color="#39c0ff", lw=2.4)
    ax.plot(ev[0], ev[2], "o", color="#39ff9e", ms=7); ax.plot(tv[0], tv[2], "o", color="#5db1ff", ms=7)
    ax.set_title("Critical structure shaded by clearance to trajectory", color="#9fd0ff", fontsize=12)
    ax.axis("off")
    cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.03)
    cb.set_label("clearance to trajectory (mm)", color="#cfe2ff")
    cb.ax.yaxis.set_tick_params(color="#cfe2ff"); plt.setp(plt.getp(cb.ax, "yticklabels"), color="#cfe2ff")
    fig.text(0.5, 0.03, "Red = trajectory grazes the structure · grey = safely distant.  "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.05, 1, 1])
    out = os.path.join(docs, "clearance_coloring.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nCLEARANCE COLORING PROOF OK")


if __name__ == "__main__":
    main()
