"""Tests for the surgical-procedure modules: tumour/GBM, ventricle/EVD, SEEG multi,
and the procedures registry."""
import numpy as np
from sulcora.volume import Volume
from sulcora.optimizer import RiskStructure
from sulcora import tumor as T, ventricle as V, session as S, procedures as P


# --------------------------------------------------------------------------- #
# Tumour / glioblastoma
# --------------------------------------------------------------------------- #
def _mask_sphere(center, radius, shape=(40, 40, 40)):
    ii, jj, kk = np.meshgrid(*[np.arange(s) for s in shape], indexing="ij")
    d = ((ii - center[0]) ** 2 + (jj - center[1]) ** 2 + (kk - center[2]) ** 2 <= radius ** 2)
    return Volume(d.astype(np.uint8), np.eye(4))


def _intensity_sphere(center, radius, val=800, bg=200, noise=30, shape=(40, 40, 40)):
    ii, jj, kk = np.meshgrid(*[np.arange(s) for s in shape], indexing="ij")
    rng = np.random.default_rng(0)
    d = (bg + rng.normal(0, noise, shape)).astype(np.float32)
    d[(ii - center[0]) ** 2 + (jj - center[1]) ** 2 + (kk - center[2]) ** 2 <= radius ** 2] = val
    return Volume(d, np.eye(4))


def test_segment_enhancing_and_volume():
    t1c = _intensity_sphere((20, 20, 20), 6, val=800, bg=200, noise=30)   # bright core on noisy brain
    core = T.segment_enhancing(t1c, percentile=98)
    assert core.data[20, 20, 20] == 1                        # centre captured
    assert 100 < int(core.data.sum()) < 4000                 # ~the sphere, not the whole volume
    assert T.mask_volume_cc(core) > 0


def test_extent_of_resection_and_rano():
    pre = _mask_sphere((20, 20, 20), 8)                        # enhancing preop mask
    remnant = np.asarray(pre.data).astype(bool) & np.asarray(_mask_sphere((26, 20, 20), 3).data).astype(bool)
    res = Volume(remnant.astype(np.uint8), np.eye(4))
    eor = T.extent_of_resection(pre, res)
    assert eor["preop_cc"] > 0 and 0 < eor["eor_pct"] < 100
    assert eor["resected_cc"] == round(eor["preop_cc"] - eor["residual_cc"], 1)
    # RANO: complete CE (0 residual), FLAIR 20cc -> 4cc (80% reduction, <=5cc) => Class 1
    c1 = T.rano_resect_class(residual_ce_cc=0.0, preop_nce_cc=20.0, residual_nce_cc=4.0)
    assert c1["class"] == 1
    # complete CE only (FLAIR barely reduced) => Class 2
    c2 = T.rano_resect_class(0.0, 20.0, 18.0)
    assert c2["class"] == 2
    # residual CE => Class 3/4
    assert T.rano_resect_class(3.0, 20.0, 15.0)["class"] == 3
    assert T.rano_resect_class(20.0, 20.0, 20.0)["class"] == 4


def test_resection_report_complete_by_default():
    ce = _mask_sphere((20, 20, 20), 6)
    fl = _mask_sphere((20, 20, 20), 10)
    rep = T.resection_report(ce, fl)              # no residual -> complete resection
    assert rep["enhancing"]["eor_pct"] == 100.0
    assert rep["rano"]["class"] == 1              # complete CE + full FLAIR removal


# --------------------------------------------------------------------------- #
# Ventricular catheter (EVD)
# --------------------------------------------------------------------------- #
def _head_ct_with_ventricle(shape=(50, 50, 50)):
    ii, jj, kk = np.meshgrid(*[np.arange(s) for s in shape], indexing="ij")
    c = np.array(shape) / 2
    r = np.sqrt((ii - c[0]) ** 2 + (jj - c[1]) ** 2 + (kk - c[2]) ** 2)
    d = np.full(shape, -1000.0, np.float32)
    d[r <= 22] = 35.0
    d[(r > 20) & (r <= 22)] = 1000.0
    # a lateral-ventricle-like CSF region, slightly right and anterior
    d[((ii - 28) ** 2 / 6 ** 2 + (jj - 30) ** 2 / 10 ** 2 + (kk - 25) ** 2 / 5 ** 2) <= 1] = 6.0
    return Volume(d, np.eye(4))


def test_segment_ventricles_and_frontal_horn():
    vol = _head_ct_with_ventricle()
    vent = V.segment_ventricles(vol)
    assert int(vent.data.sum()) > 0
    tgt = V.frontal_horn_target(vent, side="right")
    assert tgt is not None and len(tgt) == 3


def test_plan_ventricular_catheter_reaches_ventricle():
    vol = _head_ct_with_ventricle()
    vent = V.segment_ventricles(vol)
    rep = V.plan_ventricular_catheter(vol, vent, side="right", max_angle_deg=60,
                                      n_theta=5, n_phi=12)
    assert rep is not None
    assert rep["catheter_length_mm"] > 0
    assert isinstance(rep["reaches_ventricle"], bool)
    # a given-entry path straight to the target has a finite length
    rep2 = V.plan_ventricular_catheter(vol, vent, entry=rep["target"] + np.array([0, 0, 25.0]))
    assert rep2["catheter_length_mm"] > 0


# --------------------------------------------------------------------------- #
# SEEG multi-electrode
# --------------------------------------------------------------------------- #
def test_plan_multi_places_nonconflicting_corridors():
    vol = Volume(np.zeros((31, 31, 31), np.float32), np.eye(4))
    vd = np.zeros((31, 31, 31), np.uint8); vd[5, 15, 8] = 1
    structures = [RiskStructure("vessels", Volume(vd, np.eye(4)), 3.0, 1.5)]
    targets = [(15, 15, 6), (18, 12, 6)]
    plans = S.plan_multi(vol, targets, structures, min_separation_mm=4.0,
                         max_angle_deg=45, n_theta=4, n_phi=8)
    assert len(plans) == 2 and all(p is not None for p in plans)


# --------------------------------------------------------------------------- #
# Registry
# --------------------------------------------------------------------------- #
def test_procedure_registry():
    keys = [p.key for p in P.PROCEDURES]
    assert "ich" in keys and "gbm" in keys and "evd" in keys and "seeg" in keys
    assert len(keys) == len(set(keys))
    built = P.list_procedures(status="built")
    assert all(p.status == "built" for p in built) and len(built) >= 6
    assert P.PROCEDURE_BY_KEY["gbm"].status == "in_development"
    assert P.PROCEDURE_BY_KEY["ich"].kind == "corridor"
