"""Speculative decoding, from scratch, readable.

The trick in one line: a tiny DRAFT model guesses k tokens ahead (cheap),
and the big MAIN model checks all k guesses in ONE forward pass (instead of
k passes). Wherever the guesses match what main would have said, we keep them
for free; at the first disagreement we take main's own token.

Why the output is EXACT (contract A13): causal attention means main's logits
at position p depend only on tokens <= p. Every token we emit is main's own
greedy argmax given its prefix — identical to running main alone, token by
token. The draft can only make it FASTER, never different.
"""

import torch


@torch.no_grad()
def greedy_generate(model, ids, n):
    """Baseline: main model alone, one forward pass per token (argmax)."""
    model.eval()
    ids = ids.clone()
    for _ in range(n):
        ctx = ids[:, -model.cfg.max_seq_len:]
        logits, _ = model(ctx)
        nxt = logits[:, -1, :].argmax(-1, keepdim=True)
        ids = torch.cat([ids, nxt], dim=1)
    return ids


@torch.no_grad()
def speculative_generate(main, draft, ids, n, k=6):
    """Draft proposes k tokens, main verifies them in a single forward.

    Returns (ids, stats) where stats counts proposed/accepted/rounds —
    acceptance rate is what turns into speed.
    """
    main.eval(); draft.eval()
    ids = ids.clone()
    stats = {"proposed": 0, "accepted": 0, "rounds": 0, "main_forwards": 0,
             "draft_forwards": 0}
    target = ids.shape[1] + n
    while ids.shape[1] < target:
        base = ids.shape[1]
        # --- 1) draft guesses k tokens greedily (its context is short: 128) ---
        prop = ids
        for _ in range(k):
            dctx = prop[:, -draft.cfg.max_seq_len:]
            dl, _ = draft(dctx)
            nxt = dl[:, -1, :].argmax(-1, keepdim=True)
            prop = torch.cat([prop, nxt], dim=1)
            stats["draft_forwards"] += 1
        # --- 2) main checks all k guesses in ONE forward ---
        mctx = prop[:, -main.cfg.max_seq_len:]
        ml, _ = main(mctx)
        stats["main_forwards"] += 1
        greedy = ml.argmax(-1)                    # main's choice at every position
        off = prop.shape[1] - mctx.shape[1]       # shift if context was sliced
        n_acc = 0
        for j in range(k):
            pos = base + j                        # token index under review
            choice = greedy[0, pos - 1 - off]     # what main would emit there
            if choice.item() == prop[0, pos].item():
                n_acc += 1                        # guess confirmed — free token
            else:
                # first disagreement: keep confirmed prefix + main's own token
                ids = torch.cat([ids, prop[:, base:base + n_acc],
                                 choice.view(1, 1)], dim=1)
                break
        else:
            # all k confirmed — and main's forward already tells us token k+1
            bonus = greedy[0, prop.shape[1] - 1 - off].view(1, 1)
            ids = torch.cat([ids, prop[:, base:base + k], bonus], dim=1)
        stats["proposed"] += k
        stats["accepted"] += n_acc
        stats["rounds"] += 1
    return ids[:, :target], stats                 # trim overshoot to exactly n


if __name__ == "__main__":
    # self-check on random weights: exactness must hold even for untrained models
    import sys
    from pathlib import Path
    sys.path.insert(0, str(Path(__file__).parent))
    from model import TinyLLM, ModelConfig

    torch.manual_seed(0)
    main = TinyLLM(ModelConfig(dim=64, n_layers=2, n_heads=2, max_seq_len=256))
    draft = TinyLLM(ModelConfig(dim=32, n_layers=1, n_heads=2, max_seq_len=128))
    ids = torch.randint(0, 4096, (1, 10))
    ref = greedy_generate(main, ids, 40)
    out, st = speculative_generate(main, draft, ids, 40, k=4)
    assert torch.equal(ref, out), "EXACTNESS BROKEN"
    print(f"exactness OK on random models | acceptance "
          f"{st['accepted']}/{st['proposed']} (random draft ~ chance)")