Mini-Nav/tests/test_hash_codec.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 (  # noqa: E402
    bits_tensor_to_cam_row,
    bits_tensor_to_hash_bytes,
    cam_row_to_hash_bytes,
    hash_bytes_to_bits_array,
    hash_bytes_to_cam_row,
)


WIDTH = 512


def _xnor_score(query_row: int, stored_row: int, *, width: int = WIDTH) -> int:
    mask = (1 << width) - 1
    return int((~(query_row ^ stored_row) & mask).bit_count())


def _hamming_distance(left: np.ndarray, right: np.ndarray) -> int:
    return int((left != right).sum())


def test_all_zero_hash_roundtrips_through_bytes_and_cam_row():
    bits = torch.zeros(WIDTH, dtype=torch.int32)

    hash_bytes = bits_tensor_to_hash_bytes(bits)
    cam_row = hash_bytes_to_cam_row(hash_bytes)
    roundtrip = cam_row_to_hash_bytes(cam_row)

    assert len(hash_bytes) == WIDTH // 8
    assert hash_bytes == bytes(WIDTH // 8)
    assert cam_row == 0
    assert roundtrip == hash_bytes


def test_all_one_hash_roundtrips_through_bytes_and_cam_row():
    bits = torch.ones(WIDTH, dtype=torch.int32)

    hash_bytes = bits_tensor_to_hash_bytes(bits)
    cam_row = hash_bytes_to_cam_row(hash_bytes)
    roundtrip = cam_row_to_hash_bytes(cam_row)

    assert hash_bytes == b"\xff" * (WIDTH // 8)
    assert cam_row == (1 << WIDTH) - 1
    assert roundtrip == hash_bytes


def test_first_bit_uses_packbits_high_bit_ordering():
    bits = torch.zeros(WIDTH, dtype=torch.int32)
    bits[0] = 1

    hash_bytes = bits_tensor_to_hash_bytes(bits)
    cam_row = hash_bytes_to_cam_row(hash_bytes)

    assert hash_bytes[0] == 0b10000000
    assert cam_row == 1 << (WIDTH - 1)


def test_last_bit_maps_to_least_significant_cam_row_bit():
    bits = torch.zeros(WIDTH, dtype=torch.int32)
    bits[WIDTH - 1] = 1

    hash_bytes = bits_tensor_to_hash_bytes(bits)
    cam_row = hash_bytes_to_cam_row(hash_bytes)

    assert hash_bytes[-1] == 0b00000001
    assert cam_row == 1


def test_hash_bytes_unpack_to_bits_array_with_packbits_ordering():
    hash_bytes = bytes([0b10100000]) + bytes((WIDTH // 8) - 1)

    bits = hash_bytes_to_bits_array(hash_bytes)

    assert bits.shape == (WIDTH,)
    assert bits.dtype == np.uint8
    assert bits[:4].tolist() == [1, 0, 1, 0]
    assert bits[4:].sum() == 0


def test_bits_tensor_to_cam_row_matches_bytes_conversion():
    bits = torch.zeros(WIDTH, dtype=torch.int32)
    bits[0] = 1
    bits[7] = 1
    bits[511] = 1

    hash_bytes = bits_tensor_to_hash_bytes(bits)

    assert bits_tensor_to_cam_row(bits) == hash_bytes_to_cam_row(hash_bytes)


def test_positive_threshold_accepts_bool_and_signed_hash_encodings():
    bool_bits = torch.zeros(WIDTH, dtype=torch.bool)
    bool_bits[0] = True
    signed_bits = torch.full((WIDTH,), -1, dtype=torch.int32)
    signed_bits[0] = 1

    assert bits_tensor_to_hash_bytes(bool_bits) == bits_tensor_to_hash_bytes(signed_bits)


def test_xnor_score_matches_width_minus_hamming_distance():
    zeros = torch.zeros(WIDTH, dtype=torch.int32)
    ones = torch.ones(WIDTH, dtype=torch.int32)
    one_bit_diff = torch.zeros(WIDTH, dtype=torch.int32)
    one_bit_diff[0] = 1

    zero_bytes = bits_tensor_to_hash_bytes(zeros)
    one_bytes = bits_tensor_to_hash_bytes(ones)
    one_diff_bytes = bits_tensor_to_hash_bytes(one_bit_diff)

    zero_row = hash_bytes_to_cam_row(zero_bytes)
    one_row = hash_bytes_to_cam_row(one_bytes)
    one_diff_row = hash_bytes_to_cam_row(one_diff_bytes)

    zero_bits = hash_bytes_to_bits_array(zero_bytes)
    one_bits = hash_bytes_to_bits_array(one_bytes)
    one_diff_bits = hash_bytes_to_bits_array(one_diff_bytes)

    assert _hamming_distance(zero_bits, zero_bits) == 0
    assert _xnor_score(zero_row, zero_row) == WIDTH

    assert _hamming_distance(zero_bits, one_diff_bits) == 1
    assert _xnor_score(zero_row, one_diff_row) == WIDTH - 1

    assert _hamming_distance(zero_bits, one_bits) == WIDTH
    assert _xnor_score(zero_row, one_row) == 0


@pytest.mark.parametrize(
    "bits",
    [
        torch.zeros(WIDTH - 1, dtype=torch.int32),
        torch.zeros(WIDTH + 1, dtype=torch.int32),
    ],
)
def test_bits_tensor_to_hash_bytes_rejects_wrong_width(bits: torch.Tensor):
    with pytest.raises(ValueError, match="exactly 512 values"):
        bits_tensor_to_hash_bytes(bits)


def test_hash_bytes_to_bits_array_rejects_wrong_byte_length():
    with pytest.raises(ValueError, match="exactly 64 bytes"):
        hash_bytes_to_bits_array(bytes(63))


@pytest.mark.parametrize("cam_row", [-1, 1 << WIDTH])
def test_cam_row_to_hash_bytes_rejects_out_of_range_rows(cam_row: int):
    with pytest.raises(ValueError, match="range"):
        cam_row_to_hash_bytes(cam_row)


def test_width_must_be_divisible_by_eight():
    with pytest.raises(ValueError, match="divisible by 8"):
        bits_tensor_to_hash_bytes(torch.zeros(7, dtype=torch.int32), width=7)


def test_scenegraph_package_exports_hash_codec_helpers():
    from scenegraph import (  # noqa: PLC0415
        bits_tensor_to_cam_row as exported_bits_tensor_to_cam_row,
        bits_tensor_to_hash_bytes as exported_bits_tensor_to_hash_bytes,
        cam_row_to_hash_bytes as exported_cam_row_to_hash_bytes,
        hash_bytes_to_bits_array as exported_hash_bytes_to_bits_array,
        hash_bytes_to_cam_row as exported_hash_bytes_to_cam_row,
    )

    assert exported_bits_tensor_to_cam_row is bits_tensor_to_cam_row
    assert exported_bits_tensor_to_hash_bytes is bits_tensor_to_hash_bytes
    assert exported_cam_row_to_hash_bytes is cam_row_to_hash_bytes
    assert exported_hash_bytes_to_bits_array is hash_bytes_to_bits_array
    assert exported_hash_bytes_to_cam_row is hash_bytes_to_cam_row