import pytest

torch = pytest.importorskip("torch")

from hippotorch.utils.ann import (_FAISS_AVAILABLE,
                                  ApproximateNearestNeighborIndex, FaissIndex)


@pytest.mark.skipif(not _FAISS_AVAILABLE, reason="FAISS not installed")
def test_faiss_matches_exact_search():
    dim = 16
    keys = torch.randn(256, dim)
    queries = torch.randn(4, dim)

    faiss_index = FaissIndex()
    faiss_index.rebuild(keys)
    scores_faiss, idx_faiss = faiss_index.search(queries, top_k=5)

    torch_index = ApproximateNearestNeighborIndex(use_faiss=False)
    torch_index.rebuild(keys)
    scores_torch, idx_torch = torch_index.search(queries, top_k=5)

    assert torch.allclose(scores_faiss, scores_torch, atol=1e-5)
    assert torch.equal(idx_faiss, idx_torch)
