"""Tests for the patient-specific segmentation module.

These exercise the registration-free *downstream* logic — turning a label map
into structure masks, scoring trajectory clearance, and overlap metrics — using
synthetic label volumes. No deep-learning model or network access is required.
"""
import numpy as np
import pytest

from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora import segmentation as S


def _affine(spacing=(1.0, 1.0, 1.0), origin=(0, 0, 0)):
    aff = np.eye(4)
    aff[0, 0], aff[1, 1], aff[2, 2] = spacing
    aff[:3, 3] = origin
    return aff


def make_segmentation(shape=(21, 21, 21)):
    """Synthetic FreeSurfer-style label map: a thalamus blob and a lateral
    ventricle blob, plus a precentral (motor) cortex voxel."""
    data = np.zeros(shape, np.int32)
    data[5, 10, 10] = 10          # Left-Thalamus  -> "deep"/thalamus
    data[6, 10, 10] = 10
    data[10, 4, 10] = 4           # Left-Lateral-Ventricle -> "ventricle"
    data[15, 15, 15] = 1024       # ctx-lh-precentral -> "eloquent"/precentral_gyrus
    vol = Volume(data, _affine())
    return S.SegmentationResult(vol, dict(S.FREESURFER_LABELS), backend="synthseg")


# --------------------------------------------------------------------------- #
# Registry / taxonomy
# --------------------------------------------------------------------------- #
def test_registry_is_well_formed():
    keys = [s.key for s in S.STRUCTURES]
    assert len(keys) == len(set(keys))                      # unique keys
    for s in S.STRUCTURES:
        assert s.group in S.GROUPS
        assert len(s.patterns) >= 1
    # every group is represented
    assert {s.group for s in S.STRUCTURES} == set(S.GROUPS)


def test_selected_specs_filters():
    elo = S._selected_specs(groups="eloquent")
    assert elo and all(s.group == "eloquent" for s in elo)
    one = S._selected_specs(keys="thalamus")
    assert [s.key for s in one] == ["thalamus"]


# --------------------------------------------------------------------------- #
# Label-map -> mask
# --------------------------------------------------------------------------- #
def test_matching_ids_dict_and_list():
    labels = {10: "Left-Thalamus", 49: "Right-Thalamus", 4: "Left-Lateral-Ventricle"}
    assert sorted(S.matching_ids(labels, ("thalamus",))) == [10, 49]
    as_list = ["Background", "Other", "Left-Thalamus"]
    assert S.matching_ids(as_list, ("thalamus",)) == [2]


def test_patterns_match_totalsegmentator_names():
    """The taxonomy must match TotalSegmentator's label names, not only
    FreeSurfer/Harvard-Oxford ones."""
    ts = {1: "brainstem", 6: "caudate_nucleus", 7: "lentiform_nucleus", 16: "thalamus",
          1001: "ventricle_body_left", 1004: "ventricle_temporal_horn_right",
          1011: "third_ventricle", 1012: "fourth_ventricle"}
    def matched(key):
        return bool(S.matching_ids(ts, S.STRUCTURE_BY_KEY[key].patterns))
    assert matched("brainstem") and matched("caudate") and matched("thalamus")
    assert matched("lentiform_nucleus")
    assert matched("lateral_ventricle")          # via the ventricle_* horn/body parts
    assert matched("third_ventricle") and matched("fourth_ventricle")


def test_mask_from_labelmap_selects_right_voxels():
    seg = make_segmentation()
    m = S.mask_from_labelmap(seg.label_volume, seg.labels, ("thalamus",))
    assert int(m.data.sum()) == 2                           # two thalamus voxels
    assert m.data[5, 10, 10] == 1 and m.data[15, 15, 15] == 0
    empty = S.mask_from_labelmap(seg.label_volume, seg.labels, ("nonexistent-structure",))
    assert int(empty.data.sum()) == 0


def test_structure_mask_by_group():
    seg = make_segmentation()
    vent = seg.structure_mask(groups="ventricle")
    assert int(vent.data.sum()) == 1 and vent.data[10, 4, 10] == 1
    deep = seg.structure_mask(groups="deep")
    assert int(deep.data.sum()) == 2                        # thalamus only here
    elo = seg.structure_mask(groups="eloquent")
    assert int(elo.data.sum()) == 1 and elo.data[15, 15, 15] == 1


def test_present_structures():
    seg = make_segmentation()
    present = set(seg.present_structures())
    assert present == {"thalamus", "lateral_ventricle", "precentral_gyrus"}
    assert seg.present_structures(groups="ventricle") == ["lateral_ventricle"]


# --------------------------------------------------------------------------- #
# Overlap metrics
# --------------------------------------------------------------------------- #
def _mask(shape, voxels):
    d = np.zeros(shape, np.uint8)
    for v in voxels:
        d[v] = 1
    return Volume(d, _affine())


def test_dice_identical_disjoint_and_partial():
    a = _mask((10, 10, 10), [(i, 0, 0) for i in range(10)])      # 10 voxels
    assert S.dice(a, a) == 1.0
    b = _mask((10, 10, 10), [(i, 1, 0) for i in range(10)])      # disjoint
    assert S.dice(a, b) == 0.0
    # half overlap: a∩c = 5, |a|=|c|=10 -> 2*5/20 = 0.5
    c = _mask((10, 10, 10), [(i, 0, 0) for i in range(5)] + [(i, 2, 0) for i in range(5)])
    assert S.dice(a, c) == 0.5
    # both empty -> NaN
    e = _mask((10, 10, 10), [])
    assert np.isnan(S.dice(e, e))


def test_overlap_metrics():
    a = _mask((10, 10, 10), [(i, 0, 0) for i in range(10)])
    c = _mask((10, 10, 10), [(i, 0, 0) for i in range(5)] + [(i, 2, 0) for i in range(5)])
    m = S.overlap_metrics(a, c)
    assert m["seg_voxels"] == 10 and m["ref_voxels"] == 10
    assert m["dice"] == 0.5
    assert abs(m["jaccard"] - (5 / 15)) < 1e-9              # inter 5, union 15


# --------------------------------------------------------------------------- #
# Trajectory clearance against the segmentation
# --------------------------------------------------------------------------- #
def test_structure_risk_matches_clearance():
    seg = make_segmentation()
    # thalamus voxels are at world (5,10,10)/(6,10,10); run a line 4mm away in x
    traj = Trajectory(entry=(9, 0, 10), target=(9, 20, 10))
    rep = S.structure_risk(traj, seg, groups="deep", n=201)
    assert rep["nearest"] == "thalamus"
    assert abs(rep["clearance_mm"] - 3.0) < 0.3             # nearest voxel is (6,10,10)
    assert rep["band"] in ("HIGH", "CAUTION")
    assert "thalamus" in rep["per_structure"]


def test_structure_risk_empty_selection():
    seg = make_segmentation()
    # no amygdala labelled -> infinite clearance, LOW band
    rep = S.structure_risk(Trajectory((0, 0, 0), (0, 0, 20)), seg, keys="amygdala")
    assert rep["nearest"] is None and rep["band"] == "LOW"
    assert rep["clearance_mm"] == float("inf")


# --------------------------------------------------------------------------- #
# Comparison helper on a shared grid (atlas resampling stubbed with a second
# synthetic label map so no nilearn/network is needed)
# --------------------------------------------------------------------------- #
def test_compare_masks_on_shared_grid():
    seg = make_segmentation()
    # a pretend "atlas" label map on the same grid: thalamus overlaps in 1/2 voxels
    atlas = np.zeros((21, 21, 21), np.int32)
    atlas[5, 10, 10] = 10          # overlaps one of the two seg thalamus voxels
    atlas[7, 10, 10] = 10          # plus one the seg missed
    atlas_vol = Volume(atlas, _affine())
    seg_mask = seg.structure_mask(keys="thalamus")
    atl_mask = S.mask_from_labelmap(atlas_vol, S.FREESURFER_LABELS, ("thalamus",))
    m = S.overlap_metrics(seg_mask, atl_mask)
    assert m["seg_voxels"] == 2 and m["ref_voxels"] == 2
    assert m["dice"] == 0.5                                 # 2*1/(2+2)


# --------------------------------------------------------------------------- #
# Backend dispatch / graceful failure when the optional extra is absent
# --------------------------------------------------------------------------- #
def test_unknown_backend_raises():
    with pytest.raises(ValueError):
        S.segment_t1(make_segmentation().label_volume, backend="bogus")


def test_totalsegmentator_missing_raises_backend_unavailable():
    pytest.importorskip  # noqa: keep import-time symbol available
    try:
        import totalsegmentator  # noqa: F401
        pytest.skip("TotalSegmentator is installed; nothing to assert about its absence")
    except ImportError:
        pass
    with pytest.raises(S.BackendUnavailable):
        S.segment_totalsegmentator(make_segmentation().label_volume)


def test_synthseg_missing_command_raises_backend_unavailable():
    import shutil
    if shutil.which("mri_synthseg") is not None:
        pytest.skip("mri_synthseg is on PATH; nothing to assert about its absence")
    with pytest.raises(S.BackendUnavailable):
        S.segment_synthseg(make_segmentation().label_volume)
