"""Glioma / glioblastoma resection planning — segmentation + extent of resection.

Glioma surgery is a *volumetric* problem, not a corridor: the goal is to remove
as much tumour as safely possible and quantify what remains. Survival is now
graded by the **RANO resect** classification on residual enhancing (CE) and
non-enhancing (T2/FLAIR) tumour, with supramaximal resection (FLAIRectomy)
associated with longer survival. This module provides:

* two-compartment segmentation — the **enhancing core** (bright on post-contrast
  T1) and the **FLAIR / non-enhancing infiltration** (a classical intensity
  baseline; a learned model can slot in later);
* volumetrics (cc) and **extent of resection** vs a resection cavity;
* the **RANO resect class** with its associated survival band.

Pure NumPy + scipy; unit-tested headlessly. Not a medical device.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume


def segment_intensity(volume: Volume, roi_mask: Volume = None, percentile=98.0,
                      keep_largest=True, min_volume_cc=0.2) -> Volume:
    """Threshold a hyperintense region (tumour) at a high percentile, optionally
    within a supplied ROI, then keep the largest / large-enough components."""
    from scipy import ndimage
    d = np.asarray(volume.data, dtype=np.float64)
    region = (np.asarray(roi_mask.data) > 0) if roi_mask is not None else (d > d.max() * 0.05)
    vals = d[region]
    if vals.size == 0:
        return Volume(np.zeros(volume.shape, np.uint8), volume.affine.copy())
    thr = float(np.percentile(vals, percentile))
    m = (d >= thr) & region
    lab, n = ndimage.label(m)
    if n:
        sizes = np.bincount(lab.ravel()); sizes[0] = 0
        voxel_cc = float(np.prod(volume.spacing)) / 1000.0
        if keep_largest:
            keep = {int(np.argmax(sizes))}
        else:
            keep = {i for i in range(1, n + 1) if sizes[i] * voxel_cc >= min_volume_cc}
        m = np.isin(lab, list(keep)) if keep else np.zeros_like(m)
    return Volume(m.astype(np.uint8), volume.affine.copy())


def segment_enhancing(t1c: Volume, roi_mask=None, percentile=98.5) -> Volume:
    """Enhancing tumour core from post-contrast T1 (bright)."""
    return segment_intensity(t1c, roi_mask, percentile=percentile)


def segment_flair(flair: Volume, roi_mask=None, percentile=92.0) -> Volume:
    """FLAIR-hyperintense / non-enhancing infiltration (broader than the core)."""
    return segment_intensity(flair, roi_mask, percentile=percentile, keep_largest=True)


def mask_volume_cc(mask: Volume) -> float:
    return round(int(np.asarray(mask.data).sum()) * float(np.prod(mask.spacing)) / 1000.0, 1)


def extent_of_resection(preop: Volume, residual: Volume) -> dict:
    """Pre-op volume, residual, resected, and EOR % for one tumour compartment."""
    voxel_cc = float(np.prod(preop.spacing)) / 1000.0
    pre = int(np.asarray(preop.data).sum())
    res = int(np.asarray(residual.data).sum()) if residual is not None else 0
    eor = (pre - res) / pre if pre > 0 else 0.0
    return {"preop_cc": round(pre * voxel_cc, 1), "residual_cc": round(res * voxel_cc, 1),
            "resected_cc": round((pre - res) * voxel_cc, 1), "eor_pct": round(100 * eor, 1)}


def rano_resect_class(residual_ce_cc: float, preop_nce_cc: float, residual_nce_cc: float) -> dict:
    """RANO resect class from residual enhancing (CE) and non-enhancing (nCE) tumour.

    Class 1 (supramaximal): near-complete CE resection + >=60% nCE reduction and
    <=5 cc residual nCE. Class 2: complete CE only. Class 3: submaximal.
    Class 4: biopsy/large residual. Survival bands are from the RANO cohorts.
    """
    ce_complete = residual_ce_cc <= 1.0
    nce_reduction = (preop_nce_cc - residual_nce_cc) / preop_nce_cc if preop_nce_cc > 0 else 0.0
    if ce_complete and residual_nce_cc <= 5.0 and nce_reduction >= 0.60:
        return {"class": 1, "label": "Supramaximal (Class 1)", "median_os_months": "~24-29"}
    if ce_complete:
        return {"class": 2, "label": "Maximal CE (Class 2)", "median_os_months": "~19"}
    if residual_ce_cc <= 5.0:
        return {"class": 3, "label": "Submaximal (Class 3)", "median_os_months": "~15"}
    return {"class": 4, "label": "Biopsy / partial (Class 4)", "median_os_months": "—"}


def resection_report(preop_ce: Volume, preop_flair: Volume,
                     residual_ce: Volume = None, residual_flair: Volume = None) -> dict:
    """Full extent-of-resection report for both compartments + RANO class.

    If residual masks are omitted, assumes complete resection of that compartment.
    """
    zero_ce = residual_ce if residual_ce is not None else Volume(np.zeros(preop_ce.shape, np.uint8), preop_ce.affine)
    zero_fl = residual_flair if residual_flair is not None else Volume(np.zeros(preop_flair.shape, np.uint8), preop_flair.affine)
    ce = extent_of_resection(preop_ce, zero_ce)
    fl = extent_of_resection(preop_flair, zero_fl)
    rano = rano_resect_class(ce["residual_cc"], fl["preop_cc"], fl["residual_cc"])
    return {"enhancing": ce, "flair": fl, "rano": rano}
