"""Tests for the anisotropic, deformation- and error-aware required-margin generator.

The claim-critical property: a structure predicted to deform TOWARD the corridor
requires a LARGER standoff than one deforming AWAY (anisotropy), and a shakier
robot (larger error spec) requires more standoff."""
import numpy as np
from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora.optimizer import RiskStructure
from sulcora import margin as M


def _mask_at(center, shape=(30, 30, 40)):
    m = np.zeros(shape, np.uint8)
    x, y, z = center
    m[x - 1:x + 2, y - 1:y + 2, z - 1:z + 2] = 1
    return Volume(m, np.eye(4))


TRAJ = Trajectory(np.array([15, 15, 39.0]), np.array([15, 15, 5.0]))   # corridor along z at x=y=15
STRUCT = _mask_at((25, 15, 20))                                         # ~10 mm lateral
TOWARD = lambda p: p + np.array([-3.0, 0, 0])                           # moves toward corridor (x:25->22)
AWAY = lambda p: p + np.array([3.0, 0, 0])                              # moves away (x:25->28)


def test_anisotropy_toward_needs_more_margin_than_away():
    rt = M.required_standoff(TRAJ, STRUCT, deform=TOWARD, robot_sigma_mm=1, robot_sigma_deg=1, contact_prob=0.01)
    ra = M.required_standoff(TRAJ, STRUCT, deform=AWAY, robot_sigma_mm=1, robot_sigma_deg=1, contact_prob=0.01)
    assert rt["toward_mm"] > 0 and ra["toward_mm"] < 0          # direction detected
    assert rt["required_mm"] > ra["required_mm"]               # ANISOTROPY (the novel feature)
    assert rt["required_mm"] - ra["required_mm"] > 4.0         # ~ 2 * 3 mm displacement


def test_required_margin_grows_with_robot_error():
    lo = M.required_standoff(TRAJ, STRUCT, deform=None, robot_sigma_mm=0.5, robot_sigma_deg=0.5)
    hi = M.required_standoff(TRAJ, STRUCT, deform=None, robot_sigma_mm=2.0, robot_sigma_deg=2.0)
    assert hi["required_mm"] > lo["required_mm"]                # shakier robot -> more standoff


def test_tighter_contact_probability_needs_more_margin():
    loose = M.required_standoff(TRAJ, STRUCT, deform=None, contact_prob=0.10)
    tight = M.required_standoff(TRAJ, STRUCT, deform=None, contact_prob=0.001)
    assert tight["required_mm"] > loose["required_mm"]


def test_corridor_margins_binding_structure_and_feasibility():
    near = RiskStructure("near_vessel", _mask_at((21, 15, 20)), 2.0, 1.0)   # ~6 mm lateral
    far = RiskStructure("far_cortex", _mask_at((28, 15, 20)), 2.0, 1.0)     # ~13 mm lateral
    out = M.corridor_margins(TRAJ, [near, far], deform=TOWARD, robot_sigma_mm=1, robot_sigma_deg=1)
    assert set(out["per_structure"]) == {"near_vessel", "far_cortex"}
    assert out["binding"] == "near_vessel"                     # closest/most-at-risk binds
    assert isinstance(out["feasible"], bool)


def test_rank_prefers_larger_headroom():
    near = RiskStructure("v", _mask_at((22, 15, 20)), 2.0, 1.0)
    safe = Trajectory(np.array([15, 15, 39.0]), np.array([15, 15, 5.0]))    # far from structure laterally? same corridor
    risky = Trajectory(np.array([20, 15, 39.0]), np.array([20, 15, 5.0]))   # corridor nearer the structure (x=20 vs 22)
    ranked = M.rank_by_required_margin([risky, safe], [near], robot_sigma_mm=1, robot_sigma_deg=1)
    # the corridor with more clearance headroom should rank first
    assert ranked[0][0] is safe


# ---- new novel elements: envelope, inverse device spec, safety certificate ----
def test_standoff_envelope_is_anisotropic():
    pts, req = M.standoff_envelope(TRAJ, STRUCT, deform=TOWARD, robot_sigma_mm=1, robot_sigma_deg=1)
    pts2, req2 = M.standoff_envelope(TRAJ, STRUCT, deform=AWAY, robot_sigma_mm=1, robot_sigma_deg=1)
    assert len(req) == len(pts) > 0
    assert req.mean() > req2.mean()          # toward-deforming surface needs more standoff everywhere on average
    assert req.max() > req.min()             # varies across surface points (non-spherical / shaped)


def test_required_device_accuracy_inverse():
    near = RiskStructure("v", _mask_at((22, 15, 20)), 2.0, 1.0)
    far = RiskStructure("c", _mask_at((28, 15, 20)), 2.0, 1.0)
    out = M.required_device_accuracy(TRAJ, [near, far], deform=TOWARD, contact_prob=0.01)
    assert out["binding"] == "v"                         # closest structure sets the required accuracy
    assert out["max_lateral_sigma_mm"] >= 0
    assert "mm" in out["required_spec"]
    # tighter contact probability demands a more accurate device (smaller allowable sigma)
    tight = M.required_device_accuracy(TRAJ, [near, far], deform=TOWARD, contact_prob=0.0001)
    assert tight["max_lateral_sigma_mm"] <= out["max_lateral_sigma_mm"]


def test_safety_certificate_bound_and_pass():
    near = RiskStructure("v", _mask_at((22, 15, 20)), 2.0, 1.0)
    cert = M.safety_certificate(TRAJ, [near], deform=TOWARD, robot_sigma_mm=1, robot_sigma_deg=1, threshold=0.01)
    assert 0.0 <= cert["p_contact_any_bound"] <= 1.0
    assert isinstance(cert["pass"], bool) and cert["binding"] == "v"
    # a shakier robot raises the contact-probability bound
    shaky = M.safety_certificate(TRAJ, [near], deform=TOWARD, robot_sigma_mm=4, robot_sigma_deg=4, threshold=0.01)
    assert shaky["p_contact_any_bound"] >= cert["p_contact_any_bound"]
