"""ICH evacuation planning on a synthetic non-contrast CT.

Segments a hyperdense clot, finds its long axis, and plans the corridor that
maximises clot traversal while clearing vessels and eloquent cortex — compared
against a naive straight-down puncture to the centroid.

RESEARCH SOFTWARE — NOT FOR CLINICAL USE.
"""
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.volume import Volume
from sulcora.optimizer import RiskStructure
from sulcora import hematoma as H
from sulcora.planning import Trajectory


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.0
    ell = ((ii - c[0]) / (0.46 * ni)) ** 2 + ((jj - c[1]) / (0.44 * nj)) ** 2 + ((kk - c[2]) / (0.46 * nk)) ** 2
    ct = np.where(ell <= 1.0, 30.0, 0.0).astype(np.float32)          # brain ~30 HU
    ct += (np.abs(ell - 1.0) < 0.05) * 1000.0                        # thin skull ~1000 HU
    # an oblique lobar clot (hyperdense ~70 HU), elongated
    hc = np.array([56.0, 44.0, 50.0]); u = np.array([0.7, 0.2, -0.68]); u /= np.linalg.norm(u)
    ref = np.array([0, 1.0, 0]); e1 = np.cross(u, ref); e1 /= np.linalg.norm(e1); e2 = np.cross(u, e1)
    P = np.stack([ii - hc[0], jj - hc[1], kk - hc[2]], -1)
    a = P @ u; b = P @ e1; d = P @ e2
    clot = (a / 11.0) ** 2 + (b / 5.0) ** 2 + (d / 5.0) ** 2 <= 1.0
    ct[clot] = 70.0
    ct += np.random.default_rng(0).normal(0, 2.0, shape).astype(np.float32)
    vol = Volume(np.clip(ct, 0, None), np.eye(4)); interior = ell <= 1.0

    mask = H.segment_hematoma(vol, hu_lo=45, hu_hi=100)
    ax = H.hematoma_axes(mask)
    print(f"segmented clot: {ax['volume_cc']} cc, long-axis extent {ax['lengths'][0]:.0f} mm")

    elo = (((ii - 64) ** 2 + (jj - 48) ** 2 + (kk - 74) ** 2) <= 9 ** 2).astype(np.uint8)
    ves = np.zeros(shape, np.uint8)
    for t in np.linspace(0, 1, 90):
        p = np.array([40, 30, 40]) + t * (np.array([60, 66, 58]) - np.array([40, 30, 40]))
        ves[(ii - p[0]) ** 2 + (jj - p[1]) ** 2 + (kk - p[2]) ** 2 <= 4] = 1
    ves = (ves & interior).astype(np.uint8)
    structures = [RiskStructure("vessels", Volume(ves, np.eye(4)), 3.0, 1.5),
                  RiskStructure("eloquent", Volume(elo, np.eye(4)), 3.0, 1.0)]

    corridors, _ = H.plan_hematoma(vol, mask, structures, max_angle_deg=62, n_theta=9, n_phi=26)
    feas = [c for c in corridors if c.feasible]
    rec = feas[0] if feas else corridors[0]
    C = ax["centroid"]
    naive = Trajectory(C + np.array([0, 0, 40.0]), C)
    ntrav = H.traversal_length(naive, mask, 220)
    print(f"recommended: traversal {rec.traversal_mm:.0f} mm, clearance {rec.min_clearance_mm:.1f} mm ({rec.band}), "
          f"depth {rec.length_mm:.0f} mm")
    print(f"naive straight-down: traversal {ntrav:.0f} mm")

    band_col = {"HIGH": "#ff5d5d", "CAUTION": "#f5b041", "LOW": "#2ecc71"}
    disp = vol.window().data; cm = np.asarray(mask.data)
    re, rt, ne = rec.entry, rec.target, naive.entry

    def panel(a2, proj, x, y, title):
        a2.imshow(disp.max(proj).T, cmap="gray", origin="lower")
        for m, cmap in ((ves, "autumn"), (elo, "cool"), (cm, "spring")):
            mp = m.max(proj); a2.imshow(np.ma.masked_where(mp.T == 0, mp.T), cmap=cmap, alpha=.8, origin="lower")
        a2.plot([ne[x], rt[x]], [ne[y], rt[y]], "--", color="#ff5d5d", lw=1.6, alpha=.9)
        a2.plot([re[x], rt[x]], [re[y], rt[y]], "-", color=band_col[rec.band], lw=2.6)
        a2.set_title(title, color="#9fd0ff", fontsize=11); a2.axis("off")

    fig, axs = plt.subplots(1, 3, figsize=(13, 4.8), facecolor="#0b0f18")
    panel(axs[0], 2, 0, 1, "Axial MIP"); panel(axs[1], 1, 0, 2, "Coronal MIP"); panel(axs[2], 0, 1, 2, "Sagittal MIP")
    fig.suptitle(f"ICH evacuation — max-traversal corridor to a {ax['volume_cc']} cc clot", color="#dce6f5", fontsize=13)
    fig.text(0.5, 0.055, "Solid = recommended (traversal %.0f mm, clearance %.1f mm, %s). "
             "Dashed red = naive straight-down (traversal %.0f mm)."
             % (rec.traversal_mm, rec.min_clearance_mm, rec.band, ntrav), color="#cfe2ff", ha="center", fontsize=10)
    fig.text(0.5, 0.02, "magenta = segmented clot, orange = vessels, cyan = eloquent. "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.08, 1, 0.93])
    out = os.path.join(docs, "ich_evacuation.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nICH PROOF OK")


if __name__ == "__main__":
    import multiprocessing as mp; mp.freeze_support(); main()
