"""Tests on Lab 11's recorded answers: the 12 test emails the big model (gpt-5.4-mini, the teacher) got wrong in
05_distillation.py, with both students' answers; the DPO run's first step (train.dpo_loss); and the tries
behind 07_rft.py's pass@1 and pass@8. Run: pytest test_labels_practice.py"""
import math

from labels_practice import copied_mistakes, coverage, dpo_loss, same_answer

# (email, true answer, teacher, student trained on the teacher's labels, student trained on true labels)
SLIPS = [
    ('K0008', {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}),
    ('K0037', {"intent": "product_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}, {"intent": "product_question", "order_id": None, "sku": "TENT-4P", "urgency": "low"}, {"intent": "product_question", "order_id": None, "sku": "TENT-4P", "urgency": "low"}, {"intent": "product_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}),
    ('K0092', {"intent": "billing", "order_id": None, "sku": "JACK-R", "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": None, "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": "BOOT-W", "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": "JACK-R", "urgency": "high"}),
    ('K0111', {"intent": "damaged_item", "order_id": "A-2672", "sku": "MUG-T", "urgency": "low"}, {"intent": "damaged_item", "order_id": "A-2672", "sku": "MUG-T", "urgency": "normal"}, {"intent": "damaged_item", "order_id": "A-2672", "sku": "MUG-T", "urgency": "low"}, {"intent": "damaged_item", "order_id": "A-2672", "sku": "MUG-T", "urgency": "low"}),
    ('K0120', {"intent": "order_status", "order_id": "A-2520", "sku": "BOOT-W", "urgency": "normal"}, {"intent": "order_status", "order_id": "A-2520", "sku": "BOOT-W", "urgency": "high"}, {"intent": "order_status", "order_id": "A-2520", "sku": "BOOT-W", "urgency": "high"}, {"intent": "order_status", "order_id": "A-2520", "sku": "BOOT-W", "urgency": "normal"}),
    ('K0123', {"intent": "stock_question", "order_id": None, "sku": "MUG-T", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "MUG-T", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "MUG-T", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "MUG-T", "urgency": "normal"}),
    ('K0130', {"intent": "stock_question", "order_id": None, "sku": "PAD-UL", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "PAD-UL", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "PAD-UL", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "PAD-UL", "urgency": "normal"}),
    ('K0156', {"intent": "stock_question", "order_id": None, "sku": "PACK-45", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "PACK-45", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "PACK-45", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "PACK-45", "urgency": "normal"}),
    ('K0199', {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": "PAD-UL", "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "high"}, {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "high"}),
    ('K0238', {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "low"}, {"intent": "stock_question", "order_id": None, "sku": "TENT-4P", "urgency": "normal"}),
    ('K0246', {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "normal"}, {"intent": "billing", "order_id": None, "sku": "FILT-1", "urgency": "normal"}, {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "normal"}, {"intent": "billing", "order_id": None, "sku": "PACK-65", "urgency": "normal"}),
    ('K0250', {"intent": "billing", "order_id": "A-8007", "sku": None, "urgency": "normal"}, {"intent": "billing", "order_id": "A-8007", "sku": None, "urgency": "high"}, {"intent": "billing", "order_id": "A-8007", "sku": None, "urgency": "normal"}, {"intent": "billing", "order_id": "A-8007", "sku": None, "urgency": "normal"})
]
TRUTH = [s[1] for s in SLIPS]
TEACHER = [s[2] for s in SLIPS]
DIST = [s[3] for s in SLIPS]
GOLD = [s[4] for s in SLIPS]


def test_same_answer():
    assert same_answer(TRUTH[0], dict(TRUTH[0]))
    assert same_answer({"intent": "billing", "order_id": None, "sku": "tent-4p", "urgency": "LOW"},
                       {"intent": "billing", "order_id": None, "sku": "TENT-4P", "urgency": "low"})
    assert not same_answer(TRUTH[1], TEACHER[1])          # the frost-proof tent: normal vs low


def test_the_copying_student_copied_six_slips():
    assert copied_mistakes(TEACHER, DIST, TRUTH) == (12, 6)
    assert copied_mistakes(TEACHER, GOLD, TRUTH) == (12, 0)


def test_a_right_teacher_answer_is_not_a_slip():
    # Pad with an email the teacher got right and the student copied: that isn't a copied mistake.
    ok = TRUTH[0]
    assert copied_mistakes(TEACHER + [ok], DIST + [ok], TRUTH + [ok]) == (12, 6)


def test_dpo_loss_starts_at_log2():
    # Step 1 of the lab's DPO run logged a loss of 0.6931: nothing has moved from the starting model yet.
    assert math.isclose(dpo_loss(-5.0, -7.0, -5.0, -7.0, beta=0.1), math.log(2), abs_tol=1e-9)


def test_dpo_loss_falls_when_the_chosen_answer_gains():
    start = dpo_loss(-5.0, -7.0, -5.0, -7.0, 0.1)
    assert dpo_loss(-4.0, -8.0, -5.0, -7.0, 0.1) < start    # chosen up, rejected down
    assert dpo_loss(-6.0, -6.0, -5.0, -7.0, 0.1) > start    # the wrong way round
    # Only the change from the starting model counts, not how likely the answers were to begin with.
    assert math.isclose(dpo_loss(-4.0, -8.0, -5.0, -7.0, 0.1), dpo_loss(-14.0, -3.0, -15.0, -2.0, 0.1))


def test_too_easy_pairs_drive_the_loss_to_zero():
    # The lab's pairs were told apart at once: by step 10 the margin was about 14 and the loss about 0.0001.
    assert dpo_loss(0.0, -140.0, 0.0, 0.0, beta=0.1) < 1e-3 < dpo_loss(0.0, 0.0, 0.0, 0.0, beta=0.1)
    assert dpo_loss(0.0, -5000.0, 0.0, 0.0, beta=0.1) >= 0.0   # no overflow on a huge margin


# The tries from 07_rft.py's cold-open examples, before any practice: Ravi (a return, urgency normal) got
# ticks on tries 1 and 7; Mei (shouts URGENT, urgency normal) got 8 crosses.
RAVI = [True, False, False, False, False, False, True, False]
MEI = [False] * 8


def test_coverage_of_two_emails():
    c = coverage({"ravi": RAVI, "mei": MEI})
    assert c["pass1"] == 0.5 and c["passk"] == 0.5 and c["never"] == 1
    assert c["kept"] == {"ravi": 0}                     # practise on the first ticked try; Mei gives nothing


def test_coverage_matches_the_lab_round_one():
    # Before practice, on 500 emails: 72 right on the first try, 180 with a tick in 8 tries, 320 never.
    graded = {}
    for i in range(500):
        if i < 72:
            graded[i] = [True] + [False] * 7
        elif i < 180:
            graded[i] = [False] * 5 + [True] + [False] * 2
        else:
            graded[i] = [False] * 8
    c = coverage(graded)
    assert round(c["pass1"], 3) == 0.144 and c["passk"] == 0.36 and c["never"] == 320
    assert len(c["kept"]) == 180 and c["kept"][100] == 5