"""What one request costs, and how much sequential work it needs, as functions of token counts.

Money uses Decimal so the cents are exact. The prices in this lesson are examples that you edit:
no real price list is built in.
"""
from decimal import Decimal

PER = Decimal(1_000_000)  # prices are quoted per million tokens


def _dec(value):
    return value if isinstance(value, Decimal) else Decimal(str(value))


def request_cost(uncached_input, cached_input, output, price_input, price_cached, price_output):
    """Dollars for one request, as a Decimal.

    uncached_input are input tokens billed at price_input, cached_input are input tokens read
    from a stored prefix and billed at price_cached, output are generated tokens billed at
    price_output. Counts are whole tokens; prices are dollars per million tokens.
    """
    counts = (uncached_input, cached_input, output)
    prices = tuple(_dec(p) for p in (price_input, price_cached, price_output))
    if min(counts) < 0 or min(prices) < 0:
        raise ValueError("counts and prices cannot be negative")
    total = counts[0] * prices[0] + counts[1] * prices[1] + counts[2] * prices[2]
    return total / PER


def sequential_passes(input_tokens, output_tokens):
    """Passes through the model that cannot overlap in time.

    The whole prompt goes through in one pass, however long it is, and that pass also yields the
    first output token. Every later output token needs a pass of its own, because it depends on
    the one before it.
    """
    if input_tokens < 1 or output_tokens < 0:
        raise ValueError("need at least one input token and a non-negative output length")
    return max(1, output_tokens)


def decode_positions_read(input_tokens, output_tokens):
    """Earlier positions read in total by the decode passes, counting a pass's own token.

    Decode pass j (j = 1 to output_tokens - 1) reads the input_tokens prompt positions plus j
    generated ones. A model of memory traffic in positions, not a measurement in bytes or seconds.
    """
    if input_tokens < 1 or output_tokens < 0:
        raise ValueError("need at least one input token and a non-negative output length")
    passes = max(0, output_tokens - 1)
    return passes * input_tokens + passes * (passes + 1) // 2


def common_prefix_len(first, second):
    """How many leading tokens two prompts share. Only that stored prefix can be read from a cache."""
    n = 0
    for a, b in zip(first, second):
        if a != b:
            break
        n += 1
    return n