"""Tests for vessel segmentation (Frangi) and clearance-to-vessel."""
import numpy as np
import pytest

from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora import vessels as V


def _tube_volume(shape=(40, 40, 40), radius=2.0, value=1.0):
    """A bright cylindrical 'vessel' along +z through the volume centre."""
    data = np.zeros(shape, np.float32)
    ci, cj = shape[0] / 2, shape[1] / 2
    ii, jj = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
    disk = (ii - ci) ** 2 + (jj - cj) ** 2 <= radius ** 2
    data[disk, :] = value
    return Volume(data, np.eye(4))


def test_vesselness_enhances_a_tube():
    vol = _tube_volume()
    resp = V.vesselness(vol, scales=(1.0, 2.0, 3.0))
    r = np.asarray(resp.data)
    assert r.shape == vol.shape
    centre = r[20, 20, 20]                 # on the vessel axis
    corner = r[2, 2, 2]                     # background
    assert centre > corner
    assert np.isfinite(r).all()


def test_vessel_mask_captures_tube_and_is_binary():
    vol = _tube_volume()
    mask = V.vessel_mask(vol, scales=(1.0, 2.0, 3.0), percentile=99.0)
    assert set(np.unique(mask.data)).issubset({0, 1})
    assert int(mask.data.sum()) > 0
    # the detected vessel sits near the tube axis, not out in the corner
    on_axis = mask.data[18:23, 18:23, :].sum()
    corner = mask.data[:4, :4, :].sum()
    assert on_axis > corner


def test_segment_vessels_dispatch():
    vol = _tube_volume()
    m = V.segment_vessels(vol, method="frangi", scales=(1.0, 2.0))
    assert m.shape == vol.shape
    with pytest.raises(NotImplementedError):
        V.segment_vessels(vol, method="model")
    with pytest.raises(ValueError):
        V.segment_vessels(vol, method="bogus")


def test_clearance_to_mask_matches_known_distance():
    # single vessel voxel at world (10,10,10); line passes 5 mm away in x
    m = np.zeros((21, 21, 21), np.uint8); m[10, 10, 10] = 1
    mask = Volume(m, np.eye(4))
    traj = Trajectory(entry=(15, 0, 10), target=(15, 20, 10))
    d, p = V.clearance_to_mask(traj, mask, n=201)
    assert abs(d - 5.0) < 0.2 and p is not None


def test_clearance_empty_mask_is_inf():
    mask = Volume(np.zeros((10, 10, 10), np.uint8), np.eye(4))
    d, p = V.clearance_to_mask(Trajectory((0, 0, 0), (0, 0, 9)), mask)
    assert d == float("inf") and p is None


def test_clearance_kdtree_matches_bruteforce():
    rng = np.random.default_rng(0)
    m = (rng.random((25, 25, 25)) > 0.9).astype(np.uint8)   # ~10% voxels are vessel
    mask = Volume(m, np.eye(4))
    traj = Trajectory((3, 3, 3), (20, 22, 18))
    d_kd, _ = V.clearance_to_mask(traj, mask, n=128)
    d_bf, _ = traj.min_clearance_to_mask(mask, n=128)        # brute force reference
    assert abs(d_kd - d_bf) < 1e-6


def test_vessel_clearance_band_and_deformation_aware():
    from sulcora.deformation import ProbeDeformation
    m = np.zeros((31, 31, 31), np.uint8); m[5, 15, 15] = 1   # vessel at world (5,15,15)
    mask = Volume(m, np.eye(4))
    traj = Trajectory(entry=(9, 0, 15), target=(9, 30, 15))  # 4 mm away (static)
    rep = V.vessel_clearance(traj, mask, n=201)
    assert abs(rep["vessel_clearance_mm"] - 4.0) < 0.2
    assert rep["vessel_band"] == "CAUTION"
    # deformation pushes the vessel -> deformation-aware clearance differs
    pd = ProbeDeformation(traj, probe_radius_mm=1.5, influence_radius_mm=18.0)
    rep2 = V.vessel_clearance(traj, mask, deformation=pd, n=201)
    assert "deformed_vessel_clearance_mm" in rep2
    assert rep2["deformed_vessel_clearance_mm"] != rep2["vessel_clearance_mm"]


def test_vessel_risk_end_to_end():
    vol = _tube_volume()
    # a trajectory well away from the central tube -> some finite clearance
    traj = Trajectory(entry=(2, 2, 0), target=(2, 2, 39))
    rep = V.vessel_risk(traj, vol, scales=(1.0, 2.0, 3.0))
    assert rep["vessel_voxels"] > 0
    assert rep["vessel_clearance_mm"] > 0
    assert rep["vessel_band"] in ("HIGH", "CAUTION", "LOW")
