Mini-Nav/tests/test_software_cam_index.py

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pytest
import torch


MINI_NAV_DIR = Path(__file__).resolve().parents[1] / "mini-nav"
sys.path.insert(0, str(MINI_NAV_DIR))

from scenegraph.hash_codec import bits_tensor_to_hash_bytes, hash_bytes_to_cam_row  # noqa: E402
from scenegraph.objectnode import ObjectNode  # noqa: E402
from scenegraph.scenegraph import SimpleSceneGraph  # noqa: E402
from scenegraph.software_cam import SoftwareCamIndex, xnor_popcount_score  # noqa: E402


WIDTH = 512


def _hash_with_ones(*indices: int) -> bytes:
    bits = torch.zeros(WIDTH, dtype=torch.int32)
    for index in indices:
        bits[index] = 1
    return bits_tensor_to_hash_bytes(bits)


def _node(obj_id: str, hash_bytes: bytes) -> ObjectNode:
    return ObjectNode(
        obj_id=obj_id,
        room_id="room_a",
        position=np.array([0.0, 0.0, 0.0], dtype=np.float32),
        visual_hash=hash_bytes,
        semantic_hash=hash_bytes,
        hit_count=1,
        last_seen_frame=0,
    )


def _scene_graph_with_hashes(*items: tuple[str, bytes]) -> SimpleSceneGraph:
    graph = SimpleSceneGraph()
    for obj_id, hash_bytes in items:
        graph.objects[obj_id] = _node(obj_id, hash_bytes)
    return graph


def test_xnor_popcount_score_counts_matching_bits():
    assert xnor_popcount_score(0b00, 0b00, width=2) == 2
    assert xnor_popcount_score(0b00, 0b01, width=2) == 1
    assert xnor_popcount_score(0b00, 0b11, width=2) == 0


def test_software_cam_index_from_scene_graph_preserves_snapshot_order_and_rows():
    hash_a = _hash_with_ones(0)
    hash_b = _hash_with_ones(1)
    graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))

    index = SoftwareCamIndex.from_scene_graph(graph)

    assert index.obj_ids == ("obj_a", "obj_b")
    assert index.hashes == (hash_a, hash_b)
    assert index.rows == (
        hash_bytes_to_cam_row(hash_a),
        hash_bytes_to_cam_row(hash_b),
    )
    assert index.width == WIDTH


def test_software_cam_index_query_returns_top1_exact_match():
    hash_a = _hash_with_ones(0)
    hash_b = _hash_with_ones(3, 5)
    graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
    index = SoftwareCamIndex.from_scene_graph(graph)

    matches = index.query(hash_b)

    assert len(matches) == 1
    assert matches[0].obj_id == "obj_b"
    assert matches[0].row_index == 1
    assert matches[0].score == WIDTH
    assert matches[0].similarity == 1.0
    assert matches[0].hash_bytes == hash_b


def test_software_cam_index_query_topk_orders_by_score_descending():
    query = _hash_with_ones(0, 1, 2)
    two_bits_different = _hash_with_ones(0)
    exact = _hash_with_ones(0, 1, 2)
    one_bit_different = _hash_with_ones(0, 1)
    graph = _scene_graph_with_hashes(
        ("obj_a", two_bits_different),
        ("obj_b", exact),
        ("obj_c", one_bit_different),
    )
    index = SoftwareCamIndex.from_scene_graph(graph)

    matches = index.query(query, top_k=3)

    assert [match.obj_id for match in matches] == ["obj_b", "obj_c", "obj_a"]
    assert [match.score for match in matches] == [WIDTH, WIDTH - 1, WIDTH - 2]


def test_software_cam_index_query_tiebreaks_by_smaller_row_index():
    query = _hash_with_ones(0)
    equal_a = _hash_with_ones(1)
    equal_b = _hash_with_ones(2)
    graph = _scene_graph_with_hashes(("obj_a", equal_a), ("obj_b", equal_b))
    index = SoftwareCamIndex.from_scene_graph(graph)

    matches = index.query(query, top_k=2)

    assert [match.obj_id for match in matches] == ["obj_a", "obj_b"]
    assert matches[0].score == matches[1].score
    assert [match.row_index for match in matches] == [0, 1]


def test_software_cam_index_query_topk_larger_than_rows_returns_all_rows():
    hash_a = _hash_with_ones(0)
    hash_b = _hash_with_ones(1)
    graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
    index = SoftwareCamIndex.from_scene_graph(graph)

    matches = index.query(hash_a, top_k=99)

    assert len(matches) == 2


def test_empty_software_cam_index_can_be_built_but_not_queried():
    index = SoftwareCamIndex.from_scene_graph(SimpleSceneGraph())
    query = _hash_with_ones(0)

    assert index.obj_ids == ()
    assert index.rows == ()
    assert index.hashes == ()
    with pytest.raises(ValueError, match="empty SoftwareCamIndex"):
        index.query(query)


def test_software_cam_index_query_rejects_invalid_top_k():
    hash_a = _hash_with_ones(0)
    graph = _scene_graph_with_hashes(("obj_a", hash_a))
    index = SoftwareCamIndex.from_scene_graph(graph)

    with pytest.raises(ValueError, match="top_k"):
        index.query(hash_a, top_k=0)


def test_software_cam_index_uses_snapshot_semantics():
    hash_a = _hash_with_ones(0)
    hash_b = _hash_with_ones(1)
    graph = _scene_graph_with_hashes(("obj_a", hash_a))

    index = SoftwareCamIndex.from_scene_graph(graph)
    graph.objects["obj_b"] = _node("obj_b", hash_b)

    assert index.obj_ids == ("obj_a",)
    assert "obj_b" not in index.obj_ids


def test_simple_scene_graph_query_by_visual_hash_returns_scene_graph_match():
    hash_a = _hash_with_ones(0)
    hash_b = _hash_with_ones(4, 8)
    graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))

    matches = graph.query_by_visual_hash(hash_b, top_k=1)

    assert len(matches) == 1
    assert matches[0].obj_id == "obj_b"
    assert matches[0].node is graph.objects["obj_b"]
    assert matches[0].row_index == 1
    assert matches[0].score == WIDTH
    assert matches[0].similarity == 1.0
    assert matches[0].hash_bytes == hash_b


def test_scenegraph_package_exports_software_cam_query_api():
    from scenegraph import (  # noqa: PLC0415
        CamMatch,
        SceneGraphMatch,
        SoftwareCamIndex as ExportedSoftwareCamIndex,
        xnor_popcount_score as exported_xnor_popcount_score,
    )
    from scenegraph.scenegraph import SceneGraphMatch as DirectSceneGraphMatch  # noqa: PLC0415
    from scenegraph.software_cam import CamMatch as DirectCamMatch  # noqa: PLC0415

    assert CamMatch is DirectCamMatch
    assert SceneGraphMatch is DirectSceneGraphMatch
    assert ExportedSoftwareCamIndex is SoftwareCamIndex
    assert exported_xnor_popcount_score is xnor_popcount_score