#!/usr/bin/env python3
"""Speculative decoding simulation: draft-then-verify with rejection sampling.

Demonstrates the core mechanism: a cheap draft model proposes K tokens, the
target model verifies them in one parallel pass, and tokens are accepted until
the first rejection. This is a simplified simulation (no real models) that
shows the accept/reject logic and why acceptance rate drives speedup.
"""
from __future__ import annotations

import random
from dataclasses import dataclass


@dataclass
class SpecResult:
    accepted_tokens: list[str]
    rejected_at: int  # index of first rejection, or -1 if all accepted
    acceptance_rate: float


def simulate_speculative_step(
    draft_tokens: list[str],
    target_tokens: list[str],
    accept_probability: float,
    rng: random.Random,
) -> SpecResult:
    """
    Draft proposes draft_tokens; target verifies. A token is accepted if the
    draft matches the target's sample. Rejection sampling: accept up to the
    first mismatch, then resample that position from the target.
    """
    accepted: list[str] = []
    rejected_at = -1

    for i, (draft, target) in enumerate(zip(draft_tokens, target_tokens)):
        if rng.random() < accept_probability and draft == target:
            accepted.append(draft)
        else:
            # Rejection: resample this position from target's distribution
            accepted.append(target)
            rejected_at = i
            break

    acceptance_rate = len(accepted) / len(draft_tokens) if draft_tokens else 0.0
    return SpecResult(accepted, rejected_at, acceptance_rate)


def main() -> None:
    rng = random.Random(42)

    # Simulate a vocabulary of a few tokens
    vocab = ["the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog"]

    print("=== Speculative decoding simulation ===\n")
    print("How it works:")
    print("  1. Draft model (cheap) proposes K tokens")
    print("  2. Target model (expensive) verifies all K in ONE pass")
    print("  3. Accept until first rejection; resample that position\n")

    # Test different acceptance rates to show the speedup relationship
    for accept_prob in [0.9, 0.7, 0.4, 0.15]:
        K = 8  # speculative tokens per step
        total_tokens = 0
        total_steps = 0
        trials = 1000

        for _ in range(trials):
            # Draft proposes K tokens
            draft = [rng.choice(vocab) for _ in range(K)]
            # Target would sample these (simulate high correlation for high accept_prob)
            target = [
                t if rng.random() < accept_prob else rng.choice(vocab)
                for t in draft
            ]
            result = simulate_speculative_step(draft, target, accept_prob, rng)
            total_tokens += len(result.accepted)
            total_steps += 1

        # Tokens emitted per expensive (verify) step
        tokens_per_step = total_tokens / total_steps
        # Speedup vs baseline (1 token per step)
        speedup = tokens_per_step

        print(f"Acceptance rate: {accept_prob:.0%}")
        print(f"  Avg tokens/step: {tokens_per_step:.2f}  (baseline = 1.0)")
        print(f"  Effective speedup: {speedup:.2f}x\n")

    print("Key takeaways:")
    print("- Acceptance rate drives speedup (2-3x realistic, K+1 theoretical max)")
    print("- Draft must correlate with target (same model family)")
    print("- Output is statistically IDENTICAL to non-speculative decoding")


if __name__ == "__main__":
    main()