import numpy as np
import pytest
from sulcora.volume import Volume
from sulcora.planning import Trajectory, risk_band


def make_volume(spacing=(1.0, 1.0, 1.0), origin=(0, 0, 0), shape=(20, 20, 20)):
    aff = np.eye(4)
    aff[0, 0], aff[1, 1], aff[2, 2] = spacing
    aff[:3, 3] = origin
    return Volume(np.zeros(shape, np.float32), aff)


def test_world_voxel_roundtrip():
    v = make_volume(spacing=(0.9, 1.1, 2.0), origin=(-10, 5, 3))
    for ijk in [(0, 0, 0), (5, 7, 9), (19, 19, 19)]:
        w = v.voxel_to_world(ijk)
        back = v.world_to_voxel(w)
        assert np.allclose(back, ijk, atol=1e-9)


def test_spacing():
    v = make_volume(spacing=(0.9, 1.1, 2.0))
    assert np.allclose(v.spacing, [0.9, 1.1, 2.0])


def test_trajectory_length_and_angle():
    t = Trajectory(entry=(0, 0, 0), target=(0, 0, 50))
    assert abs(t.length - 50.0) < 1e-9
    assert abs(t.angle_from_axis((0, 0, 1)) - 0.0) < 1e-9
    t2 = Trajectory(entry=(0, 0, 0), target=(50, 0, 0))
    assert abs(t2.angle_from_axis((0, 0, 1)) - 90.0) < 1e-6


def test_trajectory_sample_endpoints():
    t = Trajectory((1, 2, 3), (4, 5, 6))
    pts = t.sample(10)
    assert np.allclose(pts[0], [1, 2, 3])
    assert np.allclose(pts[-1], [4, 5, 6])


def test_min_clearance_to_mask():
    # mask = single voxel at world (10,10,10); trajectory passes 5mm away in x
    v = make_volume(shape=(21, 21, 21))
    v.data[10, 10, 10] = 1
    traj = Trajectory(entry=(15, 0, 10), target=(15, 20, 10))  # line x=15, z=10
    d, p = traj.min_clearance_to_mask(v, n=201)
    assert abs(d - 5.0) < 0.2
    assert p is not None


def test_passes_through():
    v = make_volume(shape=(21, 21, 21))
    v.data[10, 10, 10] = 1
    hit = Trajectory((10, 0, 10), (10, 20, 10))
    miss = Trajectory((0, 0, 0), (0, 20, 0))
    assert hit.passes_through(v, n=201) is True
    assert miss.passes_through(v, n=201) is False


def test_risk_band():
    assert risk_band(1.0) == "HIGH"
    assert risk_band(4.0) == "CAUTION"
    assert risk_band(9.0) == "LOW"


def test_trilinear_sampling_constant_field():
    v = make_volume(shape=(10, 10, 10))
    v.data[:] = 7.0
    assert abs(v.sample_world((3.3, 4.7, 5.1)) - 7.0) < 1e-6


def test_atlas_structure_mask_and_risk():
    # synthetic labelled atlas (no network): label 2 = "Motor" at a known voxel
    from sulcora import atlas as A
    aff = np.eye(4)
    data = np.zeros((21, 21, 21), np.int16)
    data[5, 10, 10] = 2          # "Motor" voxel at world (5,10,10)
    av = Volume(data, aff)
    labels = ["Background", "Other", "Motor"]
    m = A.structure_mask(av, labels, ["Motor"])
    assert int(m.data.sum()) == 1
    # trajectory passing 4mm away from the motor voxel in x
    traj = Trajectory(entry=(9, 0, 10), target=(9, 20, 10))
    rep = A.structure_risk(traj, av, labels, names=["Motor"], n=201)
    assert rep["nearest"] == "Motor"
    assert abs(rep["clearance_mm"] - 4.0) < 0.2
    assert rep["band"] == "CAUTION"


def test_deformation_field_properties():
    from sulcora.deformation import ProbeDeformation
    traj = Trajectory(entry=(0, 0, 0), target=(0, 0, 60))   # probe along +z
    pd = ProbeDeformation(traj, probe_radius_mm=1.5, influence_radius_mm=18.0)
    near = np.array([[4.0, 0.0, 30.0]])
    far = np.array([[40.0, 0.0, 30.0]])                     # beyond influence
    mag_near = pd.magnitude(near)[0]
    assert mag_near > 0.0
    assert pd.magnitude(far)[0] == 0.0                      # localised
    vec = pd.field(near)[0]
    assert vec[0] > 0 and abs(vec[2]) < 1e-6                # radial, outward
    d2 = pd.magnitude(np.array([[8.0, 0.0, 30.0]]))[0]
    assert mag_near > d2 > 0.0                              # monotonic decay
    assert pd.max_shift_mm <= 1.5 + 1e-9                    # bounded by radius


def test_deformation_aware_risk():
    from sulcora import risk as R
    from sulcora.deformation import ProbeDeformation
    aff = np.eye(4)
    data = np.zeros((31, 31, 31), np.int16)
    data[5, 15, 15] = 2                                     # "Motor" at world (5,15,15)
    av = Volume(data, aff)
    labels = ["Background", "Other", "Motor"]
    # probe along y at x=9, z=15 -> 4mm from the structure (static)
    traj = Trajectory(entry=(9, 0, 15), target=(9, 30, 15))
    pd = ProbeDeformation(traj, probe_radius_mm=1.5, influence_radius_mm=18.0)
    rep = R.deformation_aware_risk(traj, av, labels, names=["Motor"], deformation=pd, n=201)
    assert abs(rep["static_clearance_mm"] - 4.0) < 0.2
    # the structure is within the influence radius, so it is pushed away from the
    # probe -> deformation-aware clearance should differ from static
    assert "deformed_clearance_mm" in rep
    assert rep["deformed_clearance_mm"] != rep["static_clearance_mm"]


def test_robotic_execution():
    from sulcora.robot import RoboticExecution
    planned = Trajectory(entry=(0, 0, 80), target=(0, 0, 0))
    # zero error -> perfect targeting
    r0 = RoboticExecution(planned, targeting_sigma_mm=0.0, angular_sigma_deg=0.0)
    s0 = r0.simulate(n=50, seed=1)
    assert s0["mean_tre_mm"] == 0.0 and s0["tre_within_tol_pct"] == 100.0
    # more error -> larger TRE, monotonic
    a = RoboticExecution(planned, targeting_sigma_mm=0.5, angular_sigma_deg=0.5).simulate(n=2000, seed=2)
    b = RoboticExecution(planned, targeting_sigma_mm=2.0, angular_sigma_deg=2.0).simulate(n=2000, seed=2)
    assert b["mean_tre_mm"] > a["mean_tre_mm"] > 0.0
    # safety scoring path works with a clearance function
    s = RoboticExecution(planned, targeting_sigma_mm=1.0).simulate(
        n=200, clearance_fn=lambda t: 8.0, clearance_min_mm=3.0, seed=3)
    assert s["pass_pct"] == s["tre_within_tol_pct"]  # clearance always 8>=3, so pass==within-tol


def test_surgical_plan_and_report(tmp_path):
    from sulcora.plan import SurgicalPlan
    from sulcora.report import build_report
    aff = np.eye(4); aff[:3, :3] *= 1.0; aff[:3, 3] = [-15, -15, -15]
    vol = Volume((np.random.default_rng(0).random((30, 30, 30)) * 200).astype(np.float32), aff)
    plan = SurgicalPlan(patient_id="CASE-001")
    plan.add("Left lead", Trajectory((10, 0, 10), (2, 0, 2)),
             risk={"static_clearance_mm": 5.0, "static_band": "CAUTION", "static_nearest": "X"})
    plan.add("Right lead", Trajectory((-10, 0, 10), (-2, 0, 2)))
    assert len(plan) == 2
    out = str(tmp_path / "plan.pdf")
    p = build_report(plan, vol, out, title="Test Plan")
    import os
    assert os.path.exists(p) and os.path.getsize(p) > 1000
