"""Procedure-appropriate target lists: ICH clots, DBS nuclei, EVD horns."""
import numpy as np
from sulcora.volume import Volume
from sulcora import hematoma as H, session as S


def _ct_two_clots(shape=(60, 60, 60)):
    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 <= 26] = 35.0
    d[(r > 24) & (r <= 26)] = 1000.0
    d[((ii - 40) ** 2 + (jj - 40) ** 2 + (kk - 32) ** 2) <= 7 ** 2] = 70.0    # clot A (~1.4 cc)
    d[((ii - 20) ** 2 + (jj - 22) ** 2 + (kk - 28) ** 2) <= 7 ** 2] = 70.0    # clot B, well separated
    return Volume(d, np.eye(4))


def test_list_hematomas_lists_clots_by_volume():
    vol = _ct_two_clots()
    clots = H.list_hematomas(vol, min_volume_cc=0.1)
    assert len(clots) >= 2
    # sorted largest-first, each has volume + centroid + location
    assert clots[0]["volume_cc"] >= clots[1]["volume_cc"]
    for c in clots:
        assert c["volume_cc"] > 0 and len(c["centroid"]) == 3 and isinstance(c["location"], str)


def test_procedure_targets_ich_are_the_clots():
    vol = _ct_two_clots()
    t = S.procedure_targets(vol, "ich")
    assert t["kind"] == "detected"
    assert len(t["items"]) >= 2 and "centroid" in t["items"][0]


def test_procedure_targets_dbs_are_named_nuclei():
    t = S.procedure_targets(Volume(np.zeros((5, 5, 5), np.float32), np.eye(4)), "dbs")
    assert t["kind"] == "named"
    names = [x["name"] for x in t["items"]]
    assert any("STN" in n for n in names) and any("GPi" in n for n in names) and any("VIM" in n for n in names)


def test_procedure_targets_evd_are_frontal_horns():
    vol = _ct_two_clots()
    # add a CSF ventricle so EVD has something to target
    d = np.asarray(vol.data).copy()
    ii, jj, kk = np.meshgrid(*[np.arange(s) for s in vol.shape], indexing="ij")
    d[((ii - 33) ** 2 / 5 ** 2 + (jj - 34) ** 2 / 8 ** 2 + (kk - 30) ** 2 / 4 ** 2) <= 1] = 6.0
    t = S.procedure_targets(Volume(d, np.eye(4)), "evd")
    assert t["kind"] == "detected"
    assert any("Frontal horn" in it["name"] for it in t["items"])


def test_procedure_targets_lesion_needs_mask_or_uses_it():
    vol = Volume(np.zeros((10, 10, 10), np.float32), np.eye(4))
    assert S.procedure_targets(vol, "biopsy")["kind"] == "needs_mask"
    m = np.zeros((10, 10, 10), np.uint8); m[4:6, 4:6, 4:6] = 1
    t = S.procedure_targets(vol, "gbm", lesion_mask=Volume(m, np.eye(4)))
    assert t["kind"] == "detected" and t["items"][0]["voxels"] == 8
