"""Vessel segmentation + clearance-to-vessel.

Segments vasculature with the multiscale Frangi vesselness filter and reports a
trajectory's minimum clearance to the vessel tree (the dominant SEEG/biopsy
safety metric), colouring the trajectory by risk band.

By default it runs on a synthetic vascular phantom (no public brain MR-angiogram
ships with the package); point it at a real CTA / MRA / TOF volume with
``--input path.nii.gz`` to run on patient data:

    python scripts/make_vessel_proof.py --input /path/to/angio.nii.gz

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.planning import Trajectory
from sulcora import vessels as V


def _add_tube(data, p0, p1, radius, value):
    p0, p1 = np.asarray(p0, float), np.asarray(p1, float)
    steps = int(np.linalg.norm(p1 - p0) * 3) + 2
    ni, nj, nk = data.shape
    for t in np.linspace(0, 1, steps):
        c = p0 + t * (p1 - p0)
        lo = np.floor(c - radius - 1).astype(int)
        hi = np.ceil(c + radius + 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 <= radius ** 2:
                        data[i, j, k] = value


def make_phantom(shape=(90, 90, 90), seed=0):
    """A faint ellipsoidal 'brain' with bright branching tubular 'vessels'."""
    rng = np.random.default_rng(seed)
    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
    # soft edge so the brain boundary is not a ridge the vesselness filter fires on
    data = (40.0 / (1.0 + np.exp((ell - 1.0) / 0.04))).astype(np.float32)
    # a small vessel tree (bright, like contrast/flow)
    trunk = (np.array([46, 20, 30]), np.array([44, 60, 62]))
    _add_tube(data, *trunk, radius=2.2, value=230)
    _add_tube(data, np.array([45, 40, 46]), np.array([66, 58, 36]), radius=1.6, value=220)
    _add_tube(data, np.array([45, 40, 46]), np.array([26, 62, 40]), radius=1.6, value=220)
    _add_tube(data, np.array([44, 60, 62]), np.array([58, 74, 70]), radius=1.3, value=210)
    data += rng.normal(0, 3.0, size=shape).astype(np.float32)
    return Volume(np.clip(data, 0, None), np.eye(4))


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    inp = None
    if "--input" in sys.argv:
        inp = sys.argv[sys.argv.index("--input") + 1]

    if inp:
        from sulcora.io import load
        vol = load(inp)
        source = f"real angiographic volume ({os.path.basename(inp)})"
    else:
        vol = make_phantom()
        source = "synthetic vascular phantom"
    print(f"input: {source}  shape={vol.shape}")

    mask = V.vessel_mask(vol, scales=(1.0, 2.0, 3.0), bright_vessels=True, percentile=99.2)
    if not inp:
        # keep only vessels inside the brain (drop any soft-edge response)
        ni, nj, nk = vol.shape
        ii, jj, kk = np.meshgrid(np.arange(ni), np.arange(nj), np.arange(nk), indexing="ij")
        c = np.array(vol.shape) / 2
        interior = (((ii - c[0]) / (0.44 * ni)) ** 2 + ((jj - c[1]) / (0.42 * nj)) ** 2
                    + ((kk - c[2]) / (0.44 * nk)) ** 2) <= 1.0
        mask = Volume((np.asarray(mask.data) & interior).astype(np.uint8), mask.affine)
    print(f"segmented vessel voxels: {int(mask.data.sum())}")

    # Vessel avoidance: fix a deep target, search entry points, pick the corridor
    # with the greatest clearance to the vessel tree (previews the auto-planner).
    ni, nj, nk = vol.shape
    c = np.array(vol.shape) / 2.0
    target = c + np.array([6.0, 8.0, -4.0])
    ax_a, ax_b, ax_c = 0.44 * ni, 0.42 * nj, 0.44 * nk
    best = None
    for th in np.linspace(0.15, 1.15, 7):            # polar (from superior)
        for ph in np.linspace(0, 2 * np.pi, 16, endpoint=False):
            entry = c + np.array([ax_a * np.sin(th) * np.cos(ph),
                                  ax_b * np.sin(th) * np.sin(ph),
                                  ax_c * np.cos(th)])
            tr = Trajectory(entry, target)
            d, _ = V.clearance_to_mask(tr, mask, n=200)
            if best is None or d > best[0]:
                best = (d, entry, tr)
    clearance, entry, traj = best
    rep = V.vessel_clearance(traj, mask, n=300)
    print(f"safest corridor found: length={traj.length:.1f}mm  "
          f"vessel clearance={rep['vessel_clearance_mm']:.1f}mm  band={rep['vessel_band']}")

    band_col = {"HIGH": "#ff5d5d", "CAUTION": "#f5b041", "LOW": "#2ecc71"}[rep["vessel_band"]]
    disp = vol.window().data
    vm = np.asarray(mask.data)

    def mip(ax, bg_axis):
        bg = disp.max(axis=bg_axis)
        ves = vm.max(axis=bg_axis)
        return bg, ves

    def panel(ax, bg, ves, line, pts, title):
        ax.imshow(bg.T, cmap="gray", origin="lower")
        m = np.ma.masked_where(ves.T == 0, ves.T)
        ax.imshow(m, cmap="autumn", alpha=0.85, origin="lower")
        ax.plot(*line, "-", color=band_col, lw=2.2)
        ax.plot(*pts[0], "o", color="#39ff9e", ms=6); ax.plot(*pts[1], "o", color="#5db1ff", ms=6)
        ax.set_title(title, color="#9fd0ff", fontsize=11); ax.axis("off")

    e, t = entry, target
    fig, ax = plt.subplots(1, 3, figsize=(13, 4.8), facecolor="#0b0f18")
    # axial MIP (project k): axes i(x),j(y)
    panel(ax[0], disp.max(2), vm.max(2), ([e[0], t[0]], [e[1], t[1]]),
          ([e[0], e[1]], [t[0], t[1]]), "Axial MIP")
    # coronal MIP (project j): axes i(x),k(z)
    panel(ax[1], disp.max(1), vm.max(1), ([e[0], t[0]], [e[2], t[2]]),
          ([e[0], e[2]], [t[0], t[2]]), "Coronal MIP")
    # sagittal MIP (project i): axes j(y),k(z)
    panel(ax[2], disp.max(0), vm.max(0), ([e[1], t[1]], [e[2], t[2]]),
          ([e[1], e[2]], [t[1], t[2]]), "Sagittal MIP")
    fig.suptitle(f"Vessel segmentation (Frangi) + vessel-avoidant corridor — "
                 f"clearance {rep['vessel_clearance_mm']:.1f} mm ({rep['vessel_band']})",
                 color="#dce6f5", fontsize=13)
    fig.text(0.5, 0.02, f"{source}. Orange = segmented vessels; trajectory coloured by "
             "vessel-clearance risk band. RESEARCH SOFTWARE — NOT FOR CLINICAL USE.",
             color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.05, 1, 0.93])
    out = os.path.join(docs, "vessel_avoidance.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nVESSEL PROOF OK")


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