from sovlm.dictionary import Dictionary
from sovlm.inverted_index import InvertedIndex
from sovlm.corpus import tokenize


def test_bm25_retrieval():
    d = Dictionary()
    inv = InvertedIndex()

    chunks = [
        "the cat sat on the mat",
        "the dog ran in the park",
        "quantum mechanics governs atomic transitions",
    ]
    for i, text in enumerate(chunks):
        ids = [d.add(tok) for tok in tokenize(text)]
        inv.add_chunk(i, text, ids)
    d.freeze()

    query = [d["the"], d["cat"]]
    results = inv.retrieve(query, top_k=2)
    assert len(results) > 0
    assert results[0][1] == 0  # chunk 0 has "cat"