"""Tests on the scores and checks Lab 05's suite recorded on the video's help center (21 questions,
gpt-5.4-nano answers judged by gpt-5.4-mini). Run: pytest test_gate_practice.py"""
from gate_practice import compare, critical_failures, cries_wolf, gate_passed, wobble

# Scores per run. Search and the answer scores are shares of the 16 answerable questions; refusals of the 5 others.
DAY_ONE = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.5625, "faithfulness": 1.0, "refusal_accuracy": 1.0}
SAME_1 = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.625, "faithfulness": 1.0, "refusal_accuracy": 1.0}
SAME_2 = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.625, "faithfulness": 1.0, "refusal_accuracy": 1.0}
SAME_3 = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.5625, "faithfulness": 1.0, "refusal_accuracy": 1.0}
SAME_4 = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.6875, "faithfulness": 1.0, "refusal_accuracy": 1.0}
SAME_5 = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.625, "faithfulness": 1.0, "refusal_accuracy": 1.0}
NO_REFUSAL_RULE = {"hit_at_1": 0.75, "recall_at_k": 0.875, "mrr_at_k": 0.8125, "citation_validity": 0.9375, "faithfulness": 0.75, "refusal_accuracy": 0.0}
ONE_PARAGRAPH = {"hit_at_1": 0.75, "recall_at_k": 0.75, "mrr_at_k": 0.75, "citation_validity": 0.75, "faithfulness": 0.875, "refusal_accuracy": 1.0}
FOUR_PARAGRAPHS = {"hit_at_1": 0.75, "recall_at_k": 0.9375, "mrr_at_k": 0.8281, "citation_validity": 0.875, "faithfulness": 1.0, "refusal_accuracy": 1.0}

# Which checks each question failed, per run (the rest passed). "critical" questions must pass every check.
DAY_ONE_FAILED = [
    ('plane', False, []),
    ('copy', False, ['citation_validity', 'hit_at_1', 'mrr_at_k']),
    ('quit', False, ['citation_validity']),
    ('secure', False, ['hit_at_1', 'mrr_at_k', 'recall_at_k']),
    ('locked', False, ['citation_validity']),
    ('nocard', False, ['citation_validity']),
    ('refund', True, ['citation_validity', 'hit_at_1', 'mrr_at_k', 'recall_at_k']),
    ('team14', False, ['citation_validity', 'hit_at_1', 'mrr_at_k']),
    ('slack', False, []),
    ('okta', False, []),
    ('vat', False, []),
    ('csv', False, []),
    ('reminders', False, []),
    ('archive', False, ['citation_validity']),
    ('reactivate', False, []),
    ('invite', False, []),
    ('student', True, []),
    ('darkmode', False, []),
    ('applepay', False, []),
    ('phone', False, []),
    ('linux', False, []),
]
NO_REFUSAL_RULE_FAILED = [
    ('plane', False, ['faithfulness']),
    ('copy', False, ['hit_at_1', 'mrr_at_k']),
    ('quit', False, []),
    ('secure', False, ['hit_at_1', 'mrr_at_k', 'recall_at_k']),
    ('locked', False, []),
    ('nocard', False, ['faithfulness']),
    ('refund', True, ['citation_validity', 'faithfulness', 'hit_at_1', 'mrr_at_k', 'recall_at_k']),
    ('team14', False, ['hit_at_1', 'mrr_at_k']),
    ('slack', False, []),
    ('okta', False, []),
    ('vat', False, []),
    ('csv', False, []),
    ('reminders', False, []),
    ('archive', False, []),
    ('reactivate', False, ['faithfulness']),
    ('invite', False, []),
    ('student', True, ['refusal_accuracy']),
    ('darkmode', False, ['refusal_accuracy']),
    ('applepay', False, ['refusal_accuracy']),
    ('phone', False, ['refusal_accuracy']),
    ('linux', False, ['refusal_accuracy']),
]
FOUR_PARAGRAPHS_FAILED = [
    ('plane', False, []),
    ('copy', False, ['hit_at_1', 'mrr_at_k']),
    ('quit', False, []),
    ('secure', False, ['hit_at_1', 'mrr_at_k', 'recall_at_k']),
    ('locked', False, []),
    ('nocard', False, ['citation_validity']),
    ('refund', True, ['hit_at_1', 'mrr_at_k']),
    ('team14', False, ['hit_at_1', 'mrr_at_k']),
    ('slack', False, []),
    ('okta', False, []),
    ('vat', False, []),
    ('csv', False, []),
    ('reminders', False, []),
    ('archive', False, ['citation_validity']),
    ('reactivate', False, []),
    ('invite', False, []),
    ('student', True, []),
    ('darkmode', False, []),
    ('applepay', False, []),
    ('phone', False, []),
    ('linux', False, []),
]

ALL_CHECKS = ("hit_at_1", "recall_at_k", "mrr_at_k", "citation_validity", "faithfulness", "refusal_accuracy")


def cases(failed):
    """Expand (id, critical, failed checks) rows into the case dicts the gate reads."""
    return [{"id": i, "critical": crit, "checks": {k: 0.0 if k in bad else 1.0 for k in ALL_CHECKS}}
            for i, crit, bad in failed]


def test_compare_with_tolerance():
    # The fire drill: the "say I don't know" rule deleted. Refusals 5 of 5 -> 0, the judge's score drops,
    # and the source-tag score goes UP (9 -> 15 of 16), though the assistant got worse.
    s = compare(NO_REFUSAL_RULE, DAY_ONE)
    assert s["refusal_accuracy"] == "regressed"
    assert s["faithfulness"] == "regressed"
    assert s["citation_validity"] == "improved"
    assert s["hit_at_1"] == s["recall_at_k"] == s["mrr_at_k"] == "ok"
    # One paragraph sent instead of three: search found the right one for 12 of 16, not 14. No allowance for search.
    assert compare(ONE_PARAGRAPH, DAY_ONE)["recall_at_k"] == "regressed"


def test_compare_allows_one_question_not_two():
    # Day one tagged 9 of 16; unchanged run 1 tagged 10. Saved run 1, then day one: a drop of 1 question (0.0625) is ok.
    assert compare(DAY_ONE, SAME_1)["citation_validity"] == "ok"
    # Saved the lucky run 4 (11 of 16), then run 3 (9): a drop of 2 questions (0.125). Same code, regressed.
    assert compare(SAME_3, SAME_4)["citation_validity"] == "regressed"
    assert compare(SAME_3, SAME_4, {**{k: 0.0 for k in SAME_3}, "citation_validity": 0.15})["citation_validity"] == "ok"


def test_critical_failures():
    assert critical_failures(cases(DAY_ONE_FAILED)) == ["refund"]
    # Without the rule, the student-discount question stopped getting "I don't know".
    assert critical_failures(cases(NO_REFUSAL_RULE_FAILED)) == ["refund", "student"]
    # Four paragraphs: refund gets the right answer, but its paragraph ranks 4th, so "ranked first" still fails.
    assert critical_failures(cases(FOUR_PARAGRAPHS_FAILED)) == ["refund"]
    assert dict((i, bad) for i, _, bad in FOUR_PARAGRAPHS_FAILED)["refund"] == ["hit_at_1", "mrr_at_k"]


def test_gate_decision():
    # Day one compared with itself: nothing can drop, but a critical question failed. Blocked.
    assert gate_passed(compare(DAY_ONE, DAY_ONE), critical_failures(cases(DAY_ONE_FAILED))) is False
    assert gate_passed(compare(NO_REFUSAL_RULE, DAY_ONE), critical_failures(cases(NO_REFUSAL_RULE_FAILED))) is False
    # The fix improves the source-tag score (10 -> 14 of 16) and recall, and drops nothing...
    s = compare(FOUR_PARAGRAPHS, SAME_1)
    assert s["citation_validity"] == "improved" and s["recall_at_k"] == "improved"
    assert "regressed" not in s.values()
    # ...but it is still blocked, fairly: refund is critical.
    assert gate_passed(s, critical_failures(cases(FOUR_PARAGRAPHS_FAILED))) is False
    # With no critical failure and no drop, it passes.
    assert gate_passed(s, []) is True


def test_measure_the_wobble():
    same = [SAME_1, SAME_2, SAME_3, SAME_4, SAME_5]
    cit = [r["citation_validity"] for r in same]
    assert [round(c * 16) for c in cit] == [10, 10, 9, 11, 10]
    assert abs(wobble(cit) - 0.125) < 1e-9
    for k in ("hit_at_1", "recall_at_k", "mrr_at_k", "faithfulness", "refusal_accuracy"):
        assert wobble([r[k] for r in same]) == 0.0
    # The lab's 0.10 allowance is smaller than the wobble it has to absorb: it can block unchanged code.
    assert cries_wolf(cit, 0.10) is True
    assert cries_wolf(cit, 0.15) is False
    assert cries_wolf([r["hit_at_1"] for r in same], 0.0) is False