"""Solution for spec_practice.py (Lab 12, speculative decoding)."""


def one_round(guesses, target):
    out = []
    for g, t in zip(guesses, target):
        if g != t:
            break
        out.append(g)
    out.append(target[len(out)])
    return out


def tokens_per_pass(alpha, k):
    if alpha >= 1:
        return k + 1
    return (1 - alpha ** (k + 1)) / (1 - alpha)


def speedup(alpha, k, cost_ratio):
    return tokens_per_pass(alpha, k) / (1 + k * cost_ratio)


def alpha_from_counts(accepted, total, k):
    passes = total - accepted
    if passes <= 0:
        return 1.0
    goal = total / passes
    lo, hi = 0.0, 1.0
    for _ in range(60):
        mid = (lo + hi) / 2
        if tokens_per_pass(mid, k) < goal:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2