"""Automatic trajectory optimizer — name a deep target, get the safest corridor.

Builds a phantom brain with a vessel tree and an eloquent region, then searches
entry corridors to a deep target and ranks them by clearance to *every* critical
structure at once. Compares the recommended corridor against a naive
straight-down approach.

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 optimizer as OPT
from sulcora.optimizer import RiskStructure


def _tube(data, p0, p1, radius, value):
    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 - 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 build_phantom(shape=(90, 90, 90), seed=0):
    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
    data = (40.0 / (1.0 + np.exp((ell - 1.0) / 0.04))).astype(np.float32)
    _tube(data, [46, 20, 30], [44, 60, 62], 2.2, 230)
    _tube(data, [45, 40, 46], [66, 58, 36], 1.6, 220)
    _tube(data, [45, 40, 46], [26, 62, 40], 1.6, 220)
    data += rng.normal(0, 3.0, size=shape).astype(np.float32)
    vol = Volume(np.clip(data, 0, None), np.eye(4))
    interior = ell <= 1.0
    # eloquent region (e.g. motor strip) as a ball near the superior cortex
    elo = (((ii - 60) ** 2 + (jj - 45) ** 2 + (kk - 70) ** 2) <= 8 ** 2).astype(np.uint8)
    return vol, interior, Volume(elo, np.eye(4))


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    vol, interior, elo_mask = build_phantom()
    vmask = V.vessel_mask(vol, scales=(1.0, 2.0, 3.0), bright_vessels=True, percentile=99.2)
    vmask = Volume((np.asarray(vmask.data) & interior).astype(np.uint8), vmask.affine)
    print(f"vessel voxels={int(vmask.data.sum())}  eloquent voxels={int(elo_mask.data.sum())}")

    structures = [RiskStructure("vessels", vmask, min_clearance_mm=3.0, weight=1.5),
                  RiskStructure("eloquent", elo_mask, min_clearance_mm=3.0, weight=1.0)]

    # a deep target sitting below the vessel convergence: a straight-down
    # approach skewers a vessel, but an angled corridor can avoid it.
    target = np.array([45.0, 38.0, 24.0])
    plans = OPT.plan_safest(target, structures, radius=36.0, axis=(0, 0, 1),
                            max_angle_deg=62, n_theta=9, n_phi=28, n=220)
    feasible = [p for p in plans if p.feasible]
    rec = plans[0]
    print(f"\nsearched {len(plans)} corridors; {len(feasible)} feasible")
    print("recommended: clearance=%.1fmm (%s) len=%.0fmm angle=%.0f°  per-structure=%s"
          % (rec.min_clearance_mm, rec.band, rec.length_mm, rec.angle_deg, rec.clearances))
    print("top alternatives:")
    for p in feasible[1:4]:
        print("  clearance=%.1fmm (%s) len=%.0fmm angle=%.0f°" %
              (p.min_clearance_mm, p.band, p.length_mm, p.angle_deg))

    # naive straight-down approach for comparison
    naive = OPT.score_trajectory(target + np.array([0, 0, 38.0]), target, structures,
                                 _trees={s.name: OPT._tree(V._mask_world_points(s.mask))
                                         for s in structures})
    print("naive straight-down: clearance=%.1fmm (%s)  feasible=%s"
          % (naive.min_clearance_mm, naive.band, naive.feasible))

    # ---- figure ----
    band_col = {"HIGH": "#ff5d5d", "CAUTION": "#f5b041", "LOW": "#2ecc71"}
    disp = vol.window().data
    vm = np.asarray(vmask.data); em = np.asarray(elo_mask.data)

    def panel(ax, proj, line_r, line_n, title):
        ax.imshow(disp.max(proj).T, cmap="gray", origin="lower")
        for m, cmap in ((vm, "autumn"), (em, "cool")):
            mp = m.max(proj)
            ax.imshow(np.ma.masked_where(mp.T == 0, mp.T), cmap=cmap, alpha=0.8, origin="lower")
        ax.plot(*line_n, "--", color="#ff5d5d", lw=1.8, alpha=0.9)          # naive
        ax.plot(*line_r, "-", color=band_col[rec.band], lw=2.6)             # recommended
        ax.set_title(title, color="#9fd0ff", fontsize=11); ax.axis("off")

    def proj_pts(e, t, a, b):
        return ([e[a], t[a]], [e[b], t[b]])

    re, rt = rec.entry, rec.target
    ne = naive.entry
    fig, ax = plt.subplots(1, 3, figsize=(13, 4.8), facecolor="#0b0f18")
    panel(ax[0], 2, proj_pts(re, rt, 0, 1), proj_pts(ne, rt, 0, 1), "Axial MIP")
    panel(ax[1], 1, proj_pts(re, rt, 0, 2), proj_pts(ne, rt, 0, 2), "Coronal MIP")
    panel(ax[2], 0, proj_pts(re, rt, 1, 2), proj_pts(ne, rt, 1, 2), "Sagittal MIP")
    fig.suptitle("Automatic trajectory optimizer — safest corridor vs. naive approach",
                 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, "Orange = vessels, cyan = eloquent region. Searched %d corridors, %d feasible. "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE." % (len(plans), len(feasible)),
             color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.08, 1, 0.93])
    out = os.path.join(docs, "trajectory_optimizer.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nOPTIMIZER PROOF OK")


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