"""
Practice: find what an average hides. Lab 14's script 05, Exercise 1: price one answer from its
tokens, put each trace in a turn group, split the traces and average each group, send history the
way each version does, and compare two versions like with like.

Fill in the TODOs, then run the tests:
    pytest test_causes_practice.py        (locally)
or press "Run tests" on the chapter page.
"""

# gpt-5.4-nano list prices, USD per 1M tokens (the lab's lab_utils.PRICES).
PRICES = {"input": 0.20, "output": 1.25}

TURN_GROUPS = ["1", "2-3", "4-7", "8+"]


def answer_cost(input_tokens, output_tokens, prices=PRICES):
    """The model's bill for one answer, in dollars: tokens read at prices["input"] and tokens
    written at prices["output"], both per million tokens."""
    return (input_tokens * prices["input"] + output_tokens * prices["output"]) / 1_000_000


def turn_group(turn):
    """The lab's turn groups: 1 -> "1", 2 or 3 -> "2-3", 4 to 7 -> "4-7", 8 or more -> "8+"."""
    if turn <= 1:
        return "1"
    if turn <= 3:
        return "2-3"
    if turn <= 7:
        return "4-7"
    return "8+"


def split(traces, key):
    """Group traces (dicts) into {value: [traces]} by key(trace), keeping each group's traces in
    their original order."""
    groups = {}
    for t in traces:
        groups.setdefault(key(t), []).append(t)
    return groups


def cost_per_1000(traces):
    """Mean answer_cost of the traces' "input" and "output" tokens, times 1,000."""
    return 1000 * sum(answer_cost(t["input"], t["output"]) for t in traces) / len(traces)


def history_to_send(history, cap):
    """The earlier turns sent with a new question. cap=None sends all of them (v2); otherwise only
    the last `cap` (v1 sends 3). cap=0 sends none. Return a new list."""
    if cap is None:
        return list(history)
    return list(history[-cap:]) if cap > 0 else []


def mix_average(cost_by_group, count_by_group):
    """The average cost when each group's cost is weighted by its count: sum(cost * count) /
    sum(count), over the groups in count_by_group. Use one version's costs with the other's counts
    to compare like with like."""
    total = sum(count_by_group.values())
    return sum(cost_by_group[g] * n for g, n in count_by_group.items()) / total


def change_by_group(old, new):
    """{group: relative change from old to new}, e.g. -0.21 for 21% cheaper, for the groups in both."""
    return {g: new[g] / old[g] - 1 for g in old if g in new}