"""Best pathway for haematoma evacuation (intracerebral haemorrhage / stroke).

The same planner that handles DBS/SEEG generalises to lesion procedures: given a
segmented haematoma, Sulcora targets its centroid and searches entry corridors
for the safest path to evacuate it — avoiding vessels and eloquent cortex.

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 import vessels as V
from sulcora import session
from sulcora.optimizer import RiskStructure, score_trajectory, _tree
from sulcora.vessels import _mask_world_points


def _tube(data, p0, p1, r, val):
    p0, p1 = np.asarray(p0, float), np.asarray(p1, float); ni, nj, nk = data.shape
    for t in np.linspace(0, 1, int(np.linalg.norm(p1 - p0) * 3) + 2):
        c = p0 + t * (p1 - p0); lo = np.floor(c - r - 1).astype(int); hi = np.ceil(c + r + 1).astype(int)
        for i in range(max(lo[0], 0), min(hi[0] + 1, ni)):
            for j in range(max(lo[1], 0), min(hi[1] + 1, nj)):
                for k in range(max(lo[2], 0), min(hi[2] + 1, nk)):
                    if (i - c[0]) ** 2 + (j - c[1]) ** 2 + (k - c[2]) ** 2 <= r * r:
                        data[i, j, k] = val


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)
    # a deep haematoma (ellipsoidal clot)
    hc = np.array([54.0, 40.0, 40.0])
    hem = ((((ii - hc[0]) / 9) ** 2 + ((jj - hc[1]) / 7) ** 2 + ((kk - hc[2]) / 7) ** 2) <= 1).astype(np.uint8)
    brain = np.where(hem > 0, 150.0, brain).astype(np.float32)          # hyperdense clot
    _tube(brain, [45, 22, 32], [44, 60, 60], 2.2, 230); _tube(brain, [44, 42, 46], [64, 56, 36], 1.6, 220)
    rng = np.random.default_rng(0); brain = np.clip(brain + rng.normal(0, 3, shape), 0, None).astype(np.float32)
    vol = Volume(brain, np.eye(4)); interior = ell <= 1.0

    vmask = V.vessel_mask(vol, scales=(1.0, 2.0, 3.0), percentile=99.3)
    vmask = Volume((np.asarray(vmask.data) & interior & (hem == 0)).astype(np.uint8), vmask.affine)
    elo = (((ii - 60) ** 2 + (jj - 46) ** 2 + (kk - 70) ** 2) <= 9 ** 2).astype(np.uint8)
    structures = [RiskStructure("vessels", vmask, 3.0, 1.5),
                  RiskStructure("eloquent", Volume(elo, np.eye(4)), 3.0, 1.0)]

    plans, info = session.plan_to_lesion(vol, Volume(hem, np.eye(4)), structures,
                                         max_angle_deg=60, n_theta=9, n_phi=26, n=220)
    feasible = [p for p in plans if p.feasible]; rec = feasible[0] if feasible else plans[0]
    print(f"haematoma: {info['volume_cc']} cc, centroid {np.round(info['centroid'],1)}")
    print(f"searched {len(plans)} corridors, {len(feasible)} feasible")
    print(f"recommended: clearance {rec.min_clearance_mm:.1f}mm ({rec.band}), depth {rec.length_mm:.0f}mm, angle {rec.angle_deg:.0f}deg")
    naive = score_trajectory(info["centroid"] + np.array([0, 0, 36.0]), info["centroid"], structures,
                             _trees={s.name: _tree(_mask_world_points(s.mask)) for s in structures})
    print(f"naive straight-down: clearance {naive.min_clearance_mm:.1f}mm ({naive.band})")

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

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

    fig, ax = plt.subplots(1, 3, figsize=(13, 4.8), facecolor="#0b0f18")
    panel(ax[0], 2, 0, 1, "Axial MIP"); panel(ax[1], 1, 0, 2, "Coronal MIP"); panel(ax[2], 0, 1, 2, "Sagittal MIP")
    fig.suptitle(f"Haematoma-evacuation pathway — safest corridor to a {info['volume_cc']} cc clot",
                 color="#dce6f5", fontsize=13)
    fig.text(0.5, 0.055, "Solid = recommended corridor (%.1f mm, %s). Dashed red = naive straight-down (%.1f mm, %s)."
             % (rec.min_clearance_mm, rec.band, naive.min_clearance_mm, naive.band), color="#cfe2ff", ha="center", fontsize=10)
    fig.text(0.5, 0.02, "magenta = haematoma, 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, "hematoma_pathway.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nHEMATOMA PROOF OK")


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