"""Tests for ICH hematoma segmentation + max-traversal corridor planning."""
import numpy as np
from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora.optimizer import RiskStructure
from sulcora import hematoma as H


def _ct_with_clot(shape=(40, 40, 40)):
    """Synthetic CT: brain ~30 HU, an elongated hyperdense clot ~70 HU along +x,
    plus a tiny stray hyperdensity that must be dropped as a small component."""
    d = np.full(shape, 30.0, np.float32)
    ii, jj, kk = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), np.arange(shape[2]), indexing="ij")
    clot = (((ii - 20) / 9.0) ** 2 + ((jj - 20) / 4.0) ** 2 + ((kk - 20) / 4.0) ** 2) <= 1.0
    d[clot] = 70.0
    d[3, 3, 3] = 70.0                          # stray speck (should be removed)
    return Volume(d, np.eye(4)), clot


def test_segment_hematoma_thresholds_and_keeps_largest():
    vol, clot = _ct_with_clot()
    mask = H.segment_hematoma(vol, hu_lo=45, hu_hi=100)
    assert set(np.unique(mask.data)).issubset({0, 1})
    assert mask.data[3, 3, 3] == 0             # stray speck dropped (largest CC only)
    # overlaps the clot strongly
    overlap = int((np.asarray(mask.data) & clot.astype(np.uint8)).sum())
    assert overlap > 0.8 * int(clot.sum())
    # brain (30 HU) is not included
    assert mask.data[0, 0, 0] == 0


def test_hematoma_axes_finds_long_axis():
    vol, _ = _ct_with_clot()
    mask = H.segment_hematoma(vol)
    ax = H.hematoma_axes(mask)
    assert ax["centroid"] is not None
    assert np.allclose(ax["centroid"], [20, 20, 20], atol=1.0)
    # elongated along x -> long axis is ~[1,0,0] and its extent is the largest
    assert abs(abs(ax["long_axis"][0]) - 1.0) < 0.1
    assert ax["lengths"][0] > ax["lengths"][1] and ax["lengths"][0] > ax["lengths"][2]
    assert ax["volume_cc"] > 0


def test_traversal_length():
    vol, _ = _ct_with_clot()
    mask = H.segment_hematoma(vol)
    through = Trajectory(entry=(2, 20, 20), target=(38, 20, 20))    # straight down the long axis
    miss = Trajectory(entry=(2, 2, 2), target=(2, 38, 2))           # nowhere near the clot
    t_through = H.traversal_length(through, mask, n=400)
    assert t_through > 12.0                                         # spans most of the ~18 mm clot
    assert H.traversal_length(miss, mask, n=400) == 0.0


def test_plan_hematoma_ranks_by_traversal_and_respects_safety():
    vol, _ = _ct_with_clot()
    mask = H.segment_hematoma(vol)
    # a vessel to avoid, off to one side
    vd = np.zeros((40, 40, 40), np.uint8); vd[20, 34, 20] = 1
    structures = [RiskStructure("vessels", Volume(vd, np.eye(4)), 3.0, 1.5)]
    corridors, ax = H.plan_hematoma(vol, mask, structures, max_angle_deg=75, n_theta=6, n_phi=16)
    assert corridors and ax["volume_cc"] > 0
    best = corridors[0]
    assert best.traversal_mm > 0
    # best-first ordering: earlier corridors traverse at least as much (within feasibility tier)
    feas = [c for c in corridors if c.feasible]
    if len(feas) > 1:
        assert feas[0].traversal_mm >= feas[-1].traversal_mm
    # a feasible corridor clears the vessel by its margin
    assert all(c.min_clearance_mm >= 3.0 for c in feas)
