"""Glioblastoma resection planning — two-compartment segmentation + extent of resection.

Segments the enhancing core and the FLAIR/non-enhancing infiltration, simulates a
resection cavity, and reports the extent of resection and the RANO resect class
(supramaximal / maximal / submaximal) with its survival band.

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 tumor as T


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([54.0, 46.0, 46.0])
    rng = np.random.default_rng(0)
    core = ((ii - c[0]) ** 2 + (jj - c[1]) ** 2 + (kk - c[2]) ** 2) <= 12 ** 2          # enhancing core
    flair = ((ii - c[0]) / 22.0) ** 2 + ((jj - c[1]) / 18.0) ** 2 + ((kk - c[2]) / 18.0) ** 2 <= 1  # infiltration

    t1c = (200 + rng.normal(0, 25, shape)).astype(np.float32); t1c[core] = 850.0
    fl = (200 + rng.normal(0, 25, shape)).astype(np.float32); fl[flair] = 700.0
    t1c_v, fl_v = Volume(t1c, np.eye(4)), Volume(fl, np.eye(4))

    ce = T.segment_enhancing(t1c_v, percentile=98)
    nce = T.segment_flair(fl_v, percentile=90)
    print(f"enhancing core {T.mask_volume_cc(ce)} cc, FLAIR infiltration {T.mask_volume_cc(nce)} cc")

    # simulate a supramaximal resection: remove all enhancing + most FLAIR (cavity)
    cavity = ((ii - c[0]) / 20.0) ** 2 + ((jj - c[1]) / 16.0) ** 2 + ((kk - c[2]) / 16.0) ** 2 <= 1
    res_ce = Volume((np.asarray(ce.data).astype(bool) & ~cavity).astype(np.uint8), np.eye(4))
    res_nce = Volume((np.asarray(nce.data).astype(bool) & ~cavity).astype(np.uint8), np.eye(4))
    rep = T.resection_report(ce, nce, res_ce, res_nce)
    print("EOR enhancing:", rep["enhancing"], "\nEOR FLAIR:", rep["flair"], "\nRANO:", rep["rano"])

    z = int(round(c[2]))
    fig, ax = plt.subplots(1, 3, figsize=(13, 4.7), facecolor="#0b0f18")

    def show(a, bg, masks, title):
        a.imshow(Volume(bg, np.eye(4)).window().data[:, :, z].T, cmap="gray", origin="lower")
        for m, cmap in masks:
            mm = np.asarray(m)[:, :, z]
            a.imshow(np.ma.masked_where(mm.T == 0, mm.T), cmap=cmap, alpha=0.55, origin="lower")
        a.set_title(title, color="#9fd0ff", fontsize=11); a.axis("off")

    show(ax[0], t1c, [(nce.data, "cool"), (ce.data, "autumn")], "Pre-op: enhancing (red) + FLAIR (cyan)")
    show(ax[1], t1c, [(cavity.astype(np.uint8), "Greens")], "Resection cavity (supramaximal)")
    show(ax[2], t1c, [(res_nce.data, "cool"), (res_ce.data, "autumn")], "Residual tumour")

    r = rep["rano"]
    fig.suptitle(f"Glioblastoma resection — {r['label']}  (median OS {r['median_os_months']} mo)",
                 color="#dce6f5", fontsize=13)
    fig.text(0.5, 0.045, "Enhancing EOR %.0f%% (residual %.1f cc) · FLAIR EOR %.0f%% (residual %.1f cc)"
             % (rep["enhancing"]["eor_pct"], rep["enhancing"]["residual_cc"],
                rep["flair"]["eor_pct"], rep["flair"]["residual_cc"]), color="#cfe2ff", ha="center", fontsize=10)
    fig.text(0.5, 0.015, "Two-compartment segmentation + extent-of-resection vs the RANO resect targets. "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.07, 1, 0.93])
    out = os.path.join(docs, "gbm_resection.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nGBM PROOF OK")


if __name__ == "__main__":
    main()
