"""Tests from the retrieval video's recorded rankings: 34 help-center paragraphs, 12 questions,
each labelled with the one paragraph that answers it. Run: pytest test_retrieval_practice.py"""
from retrieval_practice import rank_of, rerank, right_first, rrf

QUESTIONS = {
    "okta": "Does it work with Okta?",
    "vat": "Where do I add our VAT number?",
    "tls": "Do you use TLS 1.3?",
    "ios": "Does it run on iOS 15?",
    "csv": "Can I upload a CSV file?",
    "locked": "Why is my account locked for 15 minutes?",
    "plane": "Will it work on a plane?",
    "copy": "How do I get a copy of everything I've made?",
    "quit": "If I quit, how long before my stuff is gone?",
    "secure": "How do I make my login more secure?",
    "mention": "How do I get a teammate's attention on a task?",
    "nocard": "Can I pay without a credit card?"
}
ANSWERS = {
    "okta": "teams_3",
    "vat": "billing_2",
    "tls": "data_2",
    "ios": "app_2",
    "csv": "data_1",
    "locked": "account_3",
    "plane": "app_0",
    "copy": "data_0",
    "quit": "cancelling_2",
    "secure": "account_1",
    "mention": "guide_5",
    "nocard": "billing_1"
}
# Full meaning (dense) rankings, and keyword (BM25) rankings of every paragraph with a matching word
DENSE = {"okta": ["teams_3", "account_1", "data_2", "account_0", "guide_6", "account_3", "app_1", "billing_0", "billing_1", "app_0", "app_2", "teams_0", "guide_5", "teams_1", "data_0", "data_1", "teams_2", "guide_7", "cancelling_3", "account_2", "app_3", "guide_8", "guide_4", "cancelling_2", "cancelling_1", "billing_3", "guide_1", "data_3", "guide_3", "cancelling_0", "guide_9", "guide_0", "guide_2", "billing_2"], "vat": ["billing_2", "billing_1", "account_2", "data_2", "account_1", "data_1", "data_0", "account_0", "billing_3", "teams_2", "cancelling_3", "cancelling_0", "teams_1", "teams_0", "billing_0", "teams_3", "guide_6", "guide_0", "guide_1", "cancelling_2", "guide_8", "data_3", "guide_5", "guide_2", "guide_3", "app_3", "cancelling_1", "guide_9", "account_3", "guide_4", "guide_7", "app_1", "app_2", "app_0"], "tls": ["data_2", "app_2", "teams_3", "billing_1", "account_1", "guide_5", "data_0", "app_0", "app_1", "account_2", "billing_0", "account_0", "cancelling_2", "data_1", "guide_6", "account_3", "app_3", "billing_3", "guide_3", "billing_2", "guide_4", "data_3", "cancelling_3", "teams_1", "cancelling_1", "guide_2", "guide_9", "teams_2", "guide_1", "guide_8", "guide_0", "cancelling_0", "guide_7", "teams_0"], "ios": ["app_2", "app_0", "app_3", "app_1", "billing_0", "data_2", "billing_1", "cancelling_1", "guide_5", "data_1", "data_0", "guide_0", "guide_9", "guide_1", "guide_3", "cancelling_2", "guide_4", "cancelling_3", "guide_7", "teams_3", "billing_3", "guide_2", "account_2", "account_3", "account_1", "guide_8", "guide_6", "cancelling_0", "teams_0", "teams_2", "billing_2", "account_0", "data_3", "teams_1"], "csv": ["data_1", "data_0", "cancelling_2", "guide_1", "billing_2", "guide_0", "billing_1", "guide_6", "data_2", "app_0", "cancelling_1", "teams_1", "teams_3", "guide_5", "guide_9", "guide_2", "guide_8", "billing_0", "account_2", "app_1", "teams_2", "teams_0", "app_2", "cancelling_0", "app_3", "guide_7", "billing_3", "data_3", "cancelling_3", "account_0", "guide_4", "account_1", "guide_3", "account_3"], "locked": ["account_3", "account_0", "cancelling_3", "account_1", "data_3", "cancelling_2", "data_2", "teams_3", "account_2", "cancelling_1", "billing_1", "billing_3", "app_1", "app_3", "app_0", "teams_1", "guide_3", "teams_0", "billing_0", "teams_2", "cancelling_0", "data_0", "guide_6", "guide_5", "guide_9", "billing_2", "guide_4", "guide_8", "guide_2", "guide_0", "guide_7", "data_1", "app_2", "guide_1"], "plane": ["app_0", "app_2", "data_2", "teams_0", "teams_3", "guide_5", "billing_0", "billing_1", "guide_4", "cancelling_1", "app_1", "teams_1", "account_1", "cancelling_3", "guide_9", "account_0", "data_1", "guide_2", "app_3", "guide_7", "data_0", "cancelling_2", "guide_1", "guide_6", "guide_8", "account_3", "guide_0", "teams_2", "cancelling_0", "account_2", "guide_3", "billing_3", "data_3", "billing_2"], "copy": ["data_0", "cancelling_2", "guide_9", "billing_2", "guide_1", "data_1", "cancelling_3", "data_3", "app_1", "cancelling_1", "guide_0", "app_0", "guide_8", "guide_4", "teams_1", "app_3", "billing_3", "billing_0", "guide_5", "account_2", "teams_2", "data_2", "guide_6", "guide_7", "guide_2", "cancelling_0", "billing_1", "account_0", "teams_3", "account_1", "teams_0", "guide_3", "account_3", "app_2"], "quit": ["data_3", "cancelling_2", "cancelling_3", "cancelling_1", "billing_3", "cancelling_0", "guide_9", "data_0", "account_3", "guide_3", "app_0", "app_1", "app_3", "teams_2", "account_2", "account_0", "teams_0", "billing_0", "data_2", "guide_0", "billing_2", "billing_1", "teams_1", "guide_4", "guide_1", "guide_6", "data_1", "guide_5", "guide_7", "guide_8", "guide_2", "account_1", "teams_3", "app_2"], "secure": ["account_3", "account_1", "data_2", "teams_3", "account_0", "account_2", "data_3", "data_0", "teams_1", "cancelling_3", "cancelling_2", "billing_1", "guide_5", "guide_4", "guide_6", "billing_0", "teams_2", "app_3", "app_0", "app_1", "guide_9", "app_2", "billing_2", "guide_0", "guide_3", "data_1", "cancelling_0", "guide_8", "teams_0", "guide_1", "cancelling_1", "guide_7", "billing_3", "guide_2"], "mention": ["guide_5", "guide_2", "guide_8", "guide_4", "teams_1", "app_1", "guide_3", "guide_7", "guide_6", "teams_2", "guide_1", "data_1", "app_3", "guide_0", "teams_0", "teams_3", "guide_9", "account_3", "account_2", "app_0", "account_0", "billing_0", "data_0", "account_1", "cancelling_3", "data_2", "cancelling_0", "billing_2", "cancelling_1", "app_2", "billing_1", "data_3", "cancelling_2", "billing_3"], "nocard": ["billing_1", "cancelling_0", "app_0", "billing_2", "cancelling_1", "billing_3", "teams_3", "billing_0", "teams_1", "cancelling_3", "account_2", "account_1", "teams_2", "account_0", "data_2", "data_3", "data_0", "app_3", "teams_0", "cancelling_2", "guide_1", "account_3", "guide_2", "guide_3", "app_2", "guide_5", "guide_4", "guide_6", "data_1", "guide_0", "guide_8", "guide_7", "guide_9", "app_1"]}
BM25 = {"okta": ["teams_3", "account_0", "cancelling_0", "guide_5", "guide_4", "teams_2", "guide_7", "guide_0", "guide_8", "guide_1", "data_2", "data_0", "guide_3", "app_1", "billing_3", "guide_9", "account_3", "guide_2", "data_3", "cancelling_3", "account_1", "account_2"], "vat": ["billing_2", "billing_3", "guide_0", "data_2", "cancelling_2", "guide_2", "guide_7", "billing_1"], "tls": ["data_2", "guide_6", "guide_8", "guide_7", "guide_0", "billing_3", "billing_1", "app_0", "cancelling_1", "account_1", "guide_1", "guide_3", "app_3", "app_1", "cancelling_3", "cancelling_0", "data_0", "cancelling_2", "guide_9", "data_3", "data_1", "guide_5", "account_0", "account_3", "billing_2", "account_2"], "ios": ["account_3", "app_2", "cancelling_0", "account_0", "teams_3", "guide_2", "account_1", "account_2", "guide_9", "app_3", "cancelling_1", "billing_1", "data_3", "cancelling_3", "guide_8", "guide_1", "guide_0", "guide_6", "guide_3", "billing_2", "billing_0"], "csv": ["data_1", "data_0", "billing_1", "guide_5", "guide_8", "app_0", "guide_1", "guide_0", "teams_1", "teams_0", "app_3", "guide_3", "cancelling_3", "cancelling_0", "guide_9", "data_3", "app_1", "billing_2", "data_2", "cancelling_2", "guide_2", "guide_7", "teams_2", "billing_3", "account_1", "billing_0", "guide_4", "guide_6", "cancelling_1", "account_0", "account_3", "account_2"], "locked": ["account_3", "account_0", "teams_0", "guide_3", "cancelling_2", "data_3", "guide_7", "data_2", "teams_3", "data_0", "billing_0", "app_3", "billing_1", "app_2", "guide_9", "cancelling_0", "account_2", "guide_1", "app_1", "app_0", "teams_1", "cancelling_3", "billing_3", "cancelling_1", "account_1", "billing_2"], "plane": ["account_0", "guide_7", "guide_5", "guide_4", "guide_2", "teams_2", "account_1", "account_3", "guide_8", "guide_1", "guide_0", "cancelling_0", "cancelling_3", "account_2", "guide_9", "guide_3", "teams_3", "app_3", "billing_0", "data_3", "data_0", "guide_6", "app_0", "teams_0", "billing_2", "billing_3", "teams_1", "app_1", "data_1", "data_2", "cancelling_2", "cancelling_1", "billing_1"], "copy": ["cancelling_2", "guide_5", "billing_3", "guide_3", "cancelling_3", "guide_7", "app_3", "guide_0", "account_3", "data_0", "guide_9", "guide_4", "billing_1", "data_3", "app_1", "guide_2", "cancelling_1", "guide_8", "guide_1", "app_0", "teams_0", "account_1", "billing_0", "guide_6", "teams_1", "data_1", "cancelling_0", "account_0", "billing_2", "account_2", "data_2"], "quit": ["cancelling_2", "guide_3", "data_1", "app_1", "account_3", "cancelling_0", "account_0", "account_1", "teams_0", "data_2", "app_2", "guide_9", "guide_7", "teams_3", "data_3", "data_0", "billing_0"], "secure": ["guide_4", "account_0", "account_3", "account_2", "guide_7", "billing_3", "billing_1", "guide_0"], "mention": ["guide_5", "guide_3", "guide_7", "guide_2", "guide_0", "billing_3", "account_3", "app_1", "guide_6", "guide_4", "data_1", "billing_1", "guide_8", "account_1", "data_0", "guide_1", "billing_0", "app_0", "teams_0", "cancelling_3", "app_3", "cancelling_0", "account_0", "billing_2", "account_2", "guide_9", "teams_1", "data_3", "data_2", "cancelling_2", "teams_3", "cancelling_1"], "nocard": ["app_0", "guide_2", "billing_1", "guide_5", "guide_8", "data_0", "guide_1", "guide_0", "teams_1", "teams_0", "app_3", "guide_3", "cancelling_3", "cancelling_0", "guide_9", "data_3", "app_1", "data_1", "billing_2", "data_2", "cancelling_2", "guide_7", "teams_2", "billing_3", "account_1", "billing_0", "guide_4", "guide_6", "cancelling_1", "account_0", "account_3", "account_2"]}

# "How do I make my login more secure?": hybrid's top 10, and the reranker model's 0-10 scores
SECURE_POOL = ["account_3", "account_0", "account_2", "guide_4", "billing_1", "account_1", "data_2", "teams_3", "guide_7", "billing_3"]
SECURE_SCORES = {"account_3": 8.0, "account_0": 2.0, "account_2": 1.0, "guide_4": 0.0, "billing_1": 0.0, "account_1": 10.0, "data_2": 7.0, "teams_3": 4.0, "guide_7": 0.0, "billing_3": 0.0}


def hybrid(q, depth=20):
    return rrf([DENSE[q], BM25[q]], depth=depth)[:10]


def test_rrf_on_two_small_lists():
    assert rrf([["a", "b", "c"], ["c", "a"]]) == ["a", "c", "b"]
    assert rrf([["a", "b"], ["b"]], depth=1) == ["a", "b"]     # only each top 1 counts


def test_scores_match_the_video():
    assert right_first(DENSE, ANSWERS) == 10
    assert right_first(BM25, ANSWERS) == 7
    assert right_first({q: hybrid(q) for q in QUESTIONS}, ANSWERS) == 9


def test_why_hybrid_lost_the_plane():
    plane, right = "plane", ANSWERS["plane"]
    assert rank_of(DENSE[plane], right) == 1
    assert rank_of(BM25[plane], right, top=40) == 23        # "works", not "work"
    assert rank_of(hybrid(plane), right) is None             # out of the merged top 10
    assert rank_of(rrf([DENSE[plane], BM25[plane]]), right, top=40) == 12
    assert rank_of(rrf([DENSE[plane], BM25[plane]], depth=None), right) == 6   # merging full lists


def test_reranker_fixes_order_not_recall():
    assert rank_of(SECURE_POOL, ANSWERS["secure"]) == 6
    assert rerank(SECURE_POOL, SECURE_SCORES)[0] == ANSWERS["secure"]
    # the plane's right paragraph was never in hybrid's pool, so no reranking can bring it back
    pool = hybrid("plane")
    assert ANSWERS["plane"] not in rerank(pool, {ANSWERS["plane"]: 10})