"""Tests on Lab 14's recorded week (script 05, Exercise 1: Thursday to Sunday outside the other planted problems,
509 answers by gpt-5.4-nano, v1 217 and v2 292). Costs per 1,000 answers include the search embedding.
Run: pytest test_causes_practice.py"""
from causes_practice import (answer_cost, change_by_group, cost_per_1000, history_to_send, mix_average, split,
                             turn_group)

# Measured per turn group: cost per 1,000 answers (USD), and how many answers each version gave.
COST = {
    "v1": {"1": 0.13018, "2-3": 0.14133, "4-7": 0.15714, "8+": 0.14628},
    "v2": {"1": 0.10316, "2-3": 0.11680, "4-7": 0.14536, "8+": 0.21031},
}
COUNT = {
    "v1": {"1": 103, "2-3": 68, "4-7": 32, "8+": 14},
    "v2": {"1": 77, "2-3": 72, "4-7": 75, "8+": 68},
}


def close(a, b, tol=0.0005):
    return abs(a - b) <= tol


def test_written_tokens_cost_more_than_read_ones():
    # A first question to v1: about 308 tokens read, 55 written.
    read, written = answer_cost(308, 0), answer_cost(0, 55)
    assert close(read, 61.6e-6, 1e-9) and close(written, 68.75e-6, 1e-9)
    assert written > read                      # 55 written tokens cost more than 308 read
    assert round(answer_cost(0, 1) / answer_cost(1, 0), 2) == 6.25
    assert close(answer_cost(1_000_000, 1_000_000, {"input": 1, "output": 2}), 3.0, 1e-9)


def test_turn_groups():
    assert [turn_group(t) for t in (1, 2, 3, 4, 7, 8, 17)] == ["1", "2-3", "2-3", "4-7", "4-7", "8+", "8+"]


def test_split_and_average_each_group():
    traces = [
        {"version": "v1", "turn": 1, "input": 300, "output": 55},
        {"version": "v2", "turn": 1, "input": 320, "output": 30},
        {"version": "v2", "turn": 9, "input": 860, "output": 31},
        {"version": "v1", "turn": 9, "input": 500, "output": 36},
        {"version": "v2", "turn": 2, "input": 400, "output": 29},
    ]
    by_version = split(traces, lambda t: t["version"])
    assert sorted(by_version) == ["v1", "v2"]
    assert [t["turn"] for t in by_version["v2"]] == [1, 9, 2]           # original order kept
    by_turn = split(by_version["v2"], lambda t: turn_group(t["turn"]))
    assert sorted(by_turn) == ["1", "2-3", "8+"]
    assert close(cost_per_1000(by_turn["8+"]), 0.21075, 1e-6)         # 860 read, 31 written
    assert close(cost_per_1000(by_version["v1"]), 0.136875, 1e-6)      # mean of two answers


def test_history_v1_caps_v2_sends_everything():
    seven = [f"turn {i}" for i in range(1, 8)]                           # the learner's 8th question
    assert history_to_send(seven, 3) == ["turn 5", "turn 6", "turn 7"]  # v1
    assert history_to_send(seven, None) == seven                         # v2: no cap
    assert history_to_send(seven, None) is not seven                     # a new list
    assert history_to_send(seven, 0) == []
    assert history_to_send(["turn 1"], 3) == ["turn 1"]


def test_the_average_shows_nothing():
    v1 = mix_average(COST["v1"], COUNT["v1"])
    v2 = mix_average(COST["v2"], COUNT["v2"])
    assert round(v1, 3) == 0.139 and round(v2, 3) == 0.142             # both about 14 cents


def test_split_by_turn_shows_both_changes():
    ch = change_by_group(COST["v1"], COST["v2"])
    assert {g: round(100 * c) for g, c in ch.items()} == {"1": -21, "2-3": -17, "4-7": -7, "8+": 44}
    assert change_by_group({"a": 2.0, "b": 1.0}, {"a": 3.0}) == {"a": 0.5}


def test_like_with_like():
    # v2's learners had more long chats by chance: 68 of 292 (23%) at turn 8 or later, against 14 of 217 (6%).
    assert round(100 * COUNT["v2"]["8+"] / sum(COUNT["v2"].values())) == 23
    on_v1_mix = mix_average(COST["v2"], COUNT["v1"])
    assert round(on_v1_mix, 3) == 0.121                                # about 12 cents, 13% below v1
    assert round(100 * (1 - on_v1_mix / mix_average(COST["v1"], COUNT["v1"]))) == 13
    # On v2's own mix, v1 would have cost more than v2 did.
    assert mix_average(COST["v1"], COUNT["v2"]) > mix_average(COST["v2"], COUNT["v2"])