csaxs_bec/tests/tests_devices/test_sample_vision_monitor.py

"""Unit tests for the Layer-2 sample vision monitor device
(csaxs_bec/devices/vision/sample_vision_monitor.py)."""

import numpy as np
import pytest
from bec_server.device_server.tests.utils import DMMock

from csaxs_bec.devices.vision import vision_toolkit as vt
from csaxs_bec.devices.vision.sample_vision_monitor import (
    SampleVisionMonitor,
    SampleVisionMonitorError,
)
from csaxs_bec.devices.vision.vision_interface import VisionInterface


class FakeCamera:
    """Minimal stand-in for a camera device -- only get_last_image(), settable per test."""

    def __init__(self, image: np.ndarray | None):
        self.image = image

    def get_last_image(self) -> np.ndarray | None:
        return self.image


def _square_image(size: int = 60, offset: int = 0) -> np.ndarray:
    return vt.make_square_test_pattern(size=size, square_side=20, offset=(offset, offset))


def _make_monitor(
    camera_names=("cam_a", "cam_b"), images=None, **monitor_kwargs
) -> SampleVisionMonitor:
    dm = DMMock()
    images = images or {}
    cameras = {name: FakeCamera(images.get(name)) for name in camera_names}
    for name, camera in cameras.items():
        dm.devices[name] = camera
    interface = VisionInterface(
        name="vision_interface_test", camera_names=list(camera_names), device_manager=dm
    )
    dm.devices["vision_interface_test"] = interface
    monitor = SampleVisionMonitor(
        name="vision_monitor_test",
        vision_interface_name="vision_interface_test",
        device_manager=dm,
        **monitor_kwargs,
    )
    dm.devices["vision_monitor_test"] = monitor
    monitor._cameras = cameras  # convenience handle for tests to mutate live frames
    return monitor


def test_missing_vision_interface_raises():
    dm = DMMock()
    monitor = SampleVisionMonitor(
        name="vision_monitor_test", vision_interface_name="does_not_exist", device_manager=dm
    )
    with pytest.raises(SampleVisionMonitorError):
        monitor.score("some_tag")


def test_score_with_no_saved_references_is_infinite_but_does_not_raise():
    monitor = _make_monitor(images={"cam_a": _square_image(), "cam_b": _square_image()})
    result = monitor.score("put_mount_start_1")
    assert result["overall_score"] == float("inf")
    for camera_summary in result["per_camera"].values():
        assert camera_summary["score"] == float("inf")
        assert camera_summary["reference_count"] == 0


def test_save_reference_then_score_against_identical_frame_scores_low():
    image = _square_image()
    monitor = _make_monitor(images={"cam_a": image, "cam_b": image})
    monitor.save_reference("put_mount_start_1")

    result = monitor.score("put_mount_start_1")

    assert result["overall_score"] < float("inf")
    for camera_summary in result["per_camera"].values():
        assert camera_summary["reference_count"] == 1
        assert camera_summary["score"] == pytest.approx(0.0, abs=1e-3)


def test_score_flags_a_camera_that_no_longer_matches_its_reference():
    matching_image = _square_image()
    monitor = _make_monitor(images={"cam_a": matching_image, "cam_b": matching_image})
    monitor.save_reference("put_mount_start_1")

    # cam_b's live view changes to something with no square at all.
    monitor._cameras["cam_b"].image = np.zeros((60, 60, 3), dtype=np.uint8)

    result = monitor.score("put_mount_start_1")

    assert result["worst_camera"] == "cam_b"
    assert result["per_camera"]["cam_b"]["score"] > result["per_camera"]["cam_a"]["score"]
    assert result["overall_score"] == result["per_camera"]["cam_b"]["score"]


def test_score_publishes_a_diagnostic_overlay_for_the_worst_camera():
    image = _square_image()
    monitor = _make_monitor(images={"cam_a": image, "cam_b": image})
    monitor.save_reference("put_mount_start_1")
    monitor.score("put_mount_start_1")

    interface = monitor._resolve_vision_interface()
    published = interface.diagnostic_image.get()
    assert published.data.shape[:2] == image.shape[:2]


def test_reference_tags_and_clear_reference():
    monitor = _make_monitor(images={"cam_a": _square_image(), "cam_b": _square_image()})
    assert monitor.reference_tags() == []

    monitor.save_reference("tag_a")
    monitor.save_reference("tag_b")
    assert set(monitor.reference_tags()) == {"tag_a", "tag_b"}

    monitor.clear_reference("tag_a")
    assert monitor.reference_tags() == ["tag_b"]

    # Clearing an unknown tag is a no-op, not an error.
    monitor.clear_reference("does_not_exist")


def test_save_reference_accumulates_multiple_references_per_tag():
    monitor = _make_monitor(images={"cam_a": _square_image(), "cam_b": _square_image()})
    monitor.save_reference("tag_a")
    monitor.save_reference("tag_a")

    result = monitor.score("tag_a")
    assert result["per_camera"]["cam_a"]["reference_count"] == 2


def test_default_warp_mode_is_euclidean():
    monitor = _make_monitor()
    assert monitor.get_warp_mode() == "euclidean"


def test_constructor_accepts_a_custom_warp_mode():
    monitor = _make_monitor(warp_mode="affine")
    assert monitor.get_warp_mode() == "affine"


def test_set_warp_mode_rejects_unknown_name():
    monitor = _make_monitor()
    with pytest.raises(ValueError):
        monitor.set_warp_mode("sepia")
    # A rejected change must not have taken effect.
    assert monitor.get_warp_mode() == "euclidean"


def test_warp_mode_affects_scoring_of_a_zoomed_reference():
    # Same shape, different scale (simulating a slightly different camera zoom) -- euclidean mode
    # (translation+rotation only) cannot compensate for this, affine mode can.
    current = vt.make_square_test_pattern(size=100, square_side=20, offset=(0, 0))
    zoomed_reference = vt.make_square_test_pattern(size=100, square_side=26, offset=(0, 0))
    monitor = _make_monitor(images={"cam_a": current, "cam_b": current}, warp_mode="euclidean")
    monitor._cameras["cam_a"].image = zoomed_reference
    monitor._cameras["cam_b"].image = zoomed_reference
    monitor.save_reference("zoom_tag")
    for camera in monitor._cameras.values():
        camera.image = current

    euclidean_score = monitor.score("zoom_tag")["overall_score"]

    monitor.set_warp_mode("affine")
    affine_score = monitor.score("zoom_tag")["overall_score"]

    assert affine_score < euclidean_score


def test_position_not_checked_by_default():
    image = _square_image()
    monitor = _make_monitor(images={"cam_a": image, "cam_b": image})
    monitor.save_reference("tag_a")

    result = monitor.score("tag_a")

    assert "position_distance_px" not in result["per_camera"]["cam_a"]


def test_position_tolerance_flags_a_camera_that_has_moved_too_far():
    reference_image = _square_image(offset=0)
    monitor = _make_monitor(
        images={"cam_a": reference_image, "cam_b": reference_image}, position_tolerance_px=5.0
    )
    monitor.save_reference("tag_a")

    # cam_b's live frame is the same shape, just shifted well past the tolerance -- the shape
    # score alone (still euclidean by default) would call this a match; position tolerance must
    # catch it independently.
    monitor._cameras["cam_b"].image = _square_image(offset=15)

    result = monitor.score("tag_a")

    assert result["per_camera"]["cam_a"]["position_within_tolerance"] is True
    assert result["per_camera"]["cam_b"]["position_within_tolerance"] is False
    assert result["per_camera"]["cam_b"]["position_distance_px"] > 5.0
    # The shape score itself must still report a match -- position and shape are independent axes.
    assert result["per_camera"]["cam_b"]["score"] == pytest.approx(0.0, abs=1e-3)


def test_get_position_tolerance_px_default_is_none():
    monitor = _make_monitor()
    assert monitor.get_position_tolerance_px() is None


def test_set_position_tolerance_px_updates_at_runtime():
    monitor = _make_monitor()
    monitor.set_position_tolerance_px(2.5)
    assert monitor.get_position_tolerance_px() == 2.5
    monitor.set_position_tolerance_px(None)
    assert monitor.get_position_tolerance_px() is None