"""
Practice: speculative decoding, in code. Lab 12's arithmetic (10_speculative.py and serving.py):
one round of draft-and-check, the tokens a round gives on average, the speed-up the formula
predicts, and the acceptance rate (alpha) worked back from a run's counts.

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


def one_round(guesses, target):
    """One round with greedy decoding. `guesses` are the draft's k tokens. `target` is what the big
    model would write at each of the k + 1 positions, given the text so far plus the guesses before it
    (all found in one pass). Keep the guesses that match `target`, in order, up to the first that
    doesn't; then add the big model's own token at the next position. Return the tokens this round
    produced (a list). With every guess right, that's k + 1 tokens; with the first one wrong, 1."""
    # TODO
    return []


def tokens_per_pass(alpha, k):
    """Expected tokens one pass of the big model produces, when each of k guesses is kept with
    probability alpha (independently) and a rejected guess throws away every guess after it.
    The big model always adds one token of its own: 1 at alpha = 0, k + 1 at alpha = 1.
    Hint: 1 + alpha + alpha**2 + ... + alpha**k."""
    # TODO
    return 0.0


def speedup(alpha, k, cost_ratio):
    """The formula's predicted speed-up over plain decoding. A round costs k draft steps (each
    `cost_ratio` of a big-model step) plus one big-model pass, which is assumed to cost one step
    no matter how many tokens it checks."""
    # TODO
    return 0.0


def alpha_from_counts(accepted, total, k):
    """Estimate alpha from a run: `accepted` of `total` output tokens came from the draft. Every pass
    gives exactly one token that isn't the draft's, so passes = total - accepted, and the tokens per
    pass were total / passes. Find the alpha in [0, 1] whose tokens_per_pass(alpha, k) equals that
    (bisection works: tokens_per_pass grows with alpha). If there were no passes (accepted == total), return 1.0."""
    # TODO
    return 0.0