"""Generation for Fugue masked-diffusion models.

The continuation starts as ``<|mask|>`` tokens after the prompt. Each model
call scores masked positions and reveals some of them, one block at a time:
the whole continuation (``block_length=None``) or LLaDA's semi-autoregressive
blocks, left to right. A call reveals, in one of three ways:

- ``random`` order with ``steps_per_block`` steps: an even share of the block
  at uniformly random positions, the reverse of the i.i.d. masking Fugue
  trains on (MDLM's sampler). Fewer steps reveal more tokens per call, each
  sampled without seeing the others.
- ``random`` order with ``draft_tokens``: exact self-speculative decoding. It
  samples what random order revealing one token per call would, but each call
  checks up to ``draft_tokens`` tokens that the previous call drafted, with
  Leviathan et al.'s (2023) rejection rule along the random order. Draft i is
  checked against row i of one batched forward, which holds drafts 0..i-1.
- ``confidence`` order: each token whose sampled probability clears
  ``threshold`` and at least the single most confident one (Fast-dLLM); with
  ``threshold=None``, the most confident even share per step.

Confidence order lets an undertrained model copy its likeliest phrase into
loops (docs/architecture.md, "Fugue"); random order keeps it on the training
distribution.

With ``block_cache``, a call runs the whole canvas and caches everything
outside the block (Fast-dLLM's DualCache, see
``FugueForMaskedDiffusion.refresh_block``); the next calls run only the block
until a row has revealed ``refresh_tokens`` more of it. Text outside the block
keeps its features from the last full pass in between. Rondo's samples matched
exact decoding within noise when refreshed every 8 tokens; refreshed only at
block starts, their perplexity under Monody rose by a fifth.

With ``cuda_graphs``, every call shape is captured once as a CUDA graph per
model and replayed, since a launch-bound eager forward costs several times
its GPU time.
"""

from __future__ import annotations

import math
import weakref
from collections import OrderedDict
from collections.abc import Callable, Hashable
from dataclasses import dataclass
from typing import Any

import torch
import torch.nn.functional as F

from tiny_gdn.fugue import BlockCache, FugueForMaskedDiffusion

DEFAULT_BLOCK_LENGTH = 32
DEFAULT_CONFIDENCE_THRESHOLD = 0.9
DEFAULT_REFRESH_TOKENS = 8
UNMASKING_ORDERS: tuple[str, ...] = ("random", "confidence")
# Eager calls before capture, so Triton has compiled and tuned every kernel.
GRAPH_WARMUP_STEPS = 2
# Larger calls are bound by GPU work rather than launches and run eagerly.
GRAPH_TOKEN_LIMIT = 32_768
GRAPH_CACHE_SIZE = 4


@dataclass(frozen=True)
class DiffusionSampling:
    new_tokens: int
    # None decodes the whole continuation as one block.
    block_length: int | None = DEFAULT_BLOCK_LENGTH
    threshold: float | None = DEFAULT_CONFIDENCE_THRESHOLD
    steps_per_block: int | None = None
    order: str = "confidence"
    temperature: float = 0.0
    top_k: int = 0
    top_p: float = 1.0
    stop_at_eos: bool = True
    block_cache: bool = False
    refresh_tokens: int = DEFAULT_REFRESH_TOKENS
    draft_tokens: int = 0
    cuda_graphs: bool = False

    def __post_init__(self) -> None:
        if self.new_tokens <= 0:
            raise ValueError("new_tokens must be positive")
        if self.block_length is not None and self.block_length <= 0:
            raise ValueError("block_length must be positive")
        if self.order not in UNMASKING_ORDERS:
            raise ValueError(f"order must be one of: {', '.join(UNMASKING_ORDERS)}")
        if self.threshold is not None and not 0.0 < self.threshold <= 1.0:
            raise ValueError("threshold must be in (0, 1]")
        if self.threshold is not None and self.order != "confidence":
            raise ValueError("a confidence threshold needs confidence order")
        if self.steps_per_block is not None and self.steps_per_block <= 0:
            raise ValueError("steps_per_block must be positive")
        if self.draft_tokens < 0:
            raise ValueError("draft_tokens cannot be negative")
        if self.draft_tokens and (self.order != "random" or self.block_cache):
            raise ValueError("drafts need random order without a block cache")
        if self.threshold is None and self.steps_per_block is None and not self.draft_tokens:
            raise ValueError("set a confidence threshold, steps_per_block or draft_tokens")
        if self.temperature < 0.0:
            raise ValueError("temperature cannot be negative")
        if self.top_k < 0 or not 0.0 < self.top_p <= 1.0:
            raise ValueError("top_k must be >= 0 and top_p in (0, 1]")
        if self.refresh_tokens <= 0:
            raise ValueError("refresh_tokens must be positive")

    @property
    def block_width(self) -> int:
        return self.block_length or self.new_tokens


@dataclass(frozen=True)
class DenoisingStep:
    """State after one model call: ``tokens`` is the whole canvas, masks included."""

    step: int
    block: int
    tokens: torch.Tensor
    revealed: int
    remaining: int


# ------------------------------------------------------------------ sampling


def _filter_logits(logits: torch.Tensor, *, top_k: int, top_p: float) -> torch.Tensor:
    if top_k > 0 and top_k < logits.shape[-1]:
        kth = torch.topk(logits, top_k, dim=-1).values[..., -1:]
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if top_p < 1.0:
        sorted_logits, order = torch.sort(logits, dim=-1, descending=True)
        cumulative = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
        drop = cumulative - sorted_logits.softmax(dim=-1) >= top_p
        sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
        logits = torch.full_like(logits, float("-inf")).scatter(-1, order, sorted_logits)
    return logits


def _distribution(logits: torch.Tensor, sampling: DiffusionSampling) -> torch.Tensor:
    """What a token is sampled from, in float64 over the last dimension; one-hot at temperature 0.

    Float32 truncates the tail and lowers the effective temperature of
    masked-diffusion samplers (Zheng et al., 2024).
    """

    if sampling.temperature == 0.0:
        return F.one_hot(logits.argmax(dim=-1), logits.shape[-1]).double()
    filtered = _filter_logits(logits.double() / sampling.temperature, top_k=sampling.top_k, top_p=sampling.top_p)
    return filtered.softmax(dim=-1)


def _propose(
    logits: torch.Tensor,
    sampling: DiffusionSampling,
    generator: torch.Generator | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """``(tokens, confidence)`` per row of ``[N, vocab]`` logits."""

    if sampling.temperature == 0.0:
        confidence, tokens = logits.float().softmax(dim=-1).max(dim=-1)
        return tokens, confidence
    probabilities = _distribution(logits, sampling)
    tokens = torch.multinomial(probabilities, 1, generator=generator).squeeze(-1)
    confidence = probabilities.gather(-1, tokens[:, None]).squeeze(-1).float()
    return tokens, confidence


def _top_ranked(priority: torch.Tensor, masked: torch.Tensor, steps_left: int) -> torch.Tensor:
    """An even share of each row's masked positions, highest ``priority`` first."""

    take = torch.ceil(masked.sum(dim=1).float() / steps_left).long()
    ranks = priority.masked_fill(~masked, -1.0).argsort(dim=1, descending=True).argsort(dim=1)
    return masked & (ranks < take[:, None])


def _reveal_confident(
    sampling: DiffusionSampling, masked: torch.Tensor, confidence: torch.Tensor, steps_left: int
) -> torch.Tensor:
    """Confidence order over ``[rows, width]``; ``confidence`` is -1 off the masks."""

    if sampling.threshold is None:
        return _top_ranked(confidence, masked, steps_left)
    reveal = masked & (confidence >= sampling.threshold)
    best = confidence.argmax(dim=1)
    rows = torch.arange(masked.shape[0], device=masked.device)
    has_masks = masked.any(dim=1)
    reveal[rows[has_masks], best[has_masks]] = True
    return reveal


def _random_order(
    batch: int, prompt_length: int, sampling: DiffusionSampling, generator: torch.Generator | None, device: torch.device
) -> torch.Tensor:
    """Each row's continuation positions ``[batch, new_tokens]`` in a uniformly random order within each block."""

    blocks = []
    for start in range(0, sampling.new_tokens, sampling.block_width):
        width = min(sampling.block_width, sampling.new_tokens - start)
        draws = torch.rand(batch, width, generator=generator, device=device)
        blocks.append(prompt_length + start + draws.argsort(dim=1))
    return torch.cat(blocks, dim=1)


def _first_eos(tokens: torch.Tensor, eos_token_id: int) -> torch.Tensor:
    """Each row's first EOS index in ``[rows, T]``, or ``T`` for rows without one."""

    length = tokens.shape[1]
    positions = torch.arange(length, device=tokens.device)
    return torch.where(tokens.eq(eos_token_id), positions, length).min(dim=1).values


def _committed_eos(tokens: torch.Tensor, mask_token_id: int, eos_token_id: int) -> torch.Tensor:
    """Rows of ``[rows, T]`` whose revealed prefix reaches an EOS."""

    length = tokens.shape[1]
    positions = torch.arange(length, device=tokens.device)
    first_mask = torch.where(tokens.eq(mask_token_id), positions, length).min(dim=1).values
    return _first_eos(tokens, eos_token_id) < first_mask


def _cut_order(
    order: torch.Tensor, start: torch.Tensor | int, stop: torch.Tensor | int, cutoff: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Move each row's pending positions ``order[start:stop]`` past ``cutoff`` behind the rest.

    Returns the order, the rest keeping its relative order, and how many positions each row moved.
    """

    index = torch.arange(order.shape[1], device=order.device)
    start = torch.as_tensor(start, device=order.device).reshape(-1, 1)
    stop = torch.as_tensor(stop, device=order.device).reshape(-1, 1)
    pending = (index >= start) & (index < stop)
    dead = pending & (order > cutoff[:, None])
    moved = dead.sum(dim=1)
    if not bool(moved.any()):
        return order, moved
    key = torch.where(index < start, 0, torch.where(pending, 1 + dead.long(), 3))
    return order.gather(1, torch.sort(key, dim=1, stable=True).indices), moved


# -------------------------------------------------------------------- graphs


class _Captured:
    """A function of tensors, captured once as a CUDA graph over static copies of its inputs."""

    def __init__(self, function: Callable[..., Any], inputs: tuple[torch.Tensor, ...], pool: Any) -> None:
        self.inputs = tuple(tensor.clone() for tensor in inputs)
        device = self.inputs[0].device
        current = torch.cuda.current_stream(device)
        stream = torch.cuda.Stream(device=device)
        stream.wait_stream(current)
        with torch.cuda.stream(stream):
            for _ in range(GRAPH_WARMUP_STEPS):
                function(*self.inputs)
        current.wait_stream(stream)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph, pool=pool, capture_error_mode="thread_local"):
            self.outputs = function(*self.inputs)

    def __call__(self, *inputs: torch.Tensor) -> Any:
        for static, tensor in zip(self.inputs, inputs, strict=True):
            static.copy_(tensor)
        self.graph.replay()
        return self.outputs


class _GraphCache:
    """One model's captured decoding graphs by call shape, least recently used dropped first."""

    def __init__(self) -> None:
        self.graphs: OrderedDict[Hashable, Any] = OrderedDict()

    def get(self, key: Hashable, build: Callable[[], Any]) -> Any:
        graph = self.graphs.pop(key, None)
        if graph is None:
            while len(self.graphs) >= GRAPH_CACHE_SIZE:
                self.graphs.popitem(last=False)
            graph = build()
        self.graphs[key] = graph
        return graph


# Graphs bake in the parameters' addresses, so they live and die with the model object.
_GRAPH_CACHES: weakref.WeakKeyDictionary[FugueForMaskedDiffusion, _GraphCache] = weakref.WeakKeyDictionary()


def _graphs(model: FugueForMaskedDiffusion) -> _GraphCache:
    cache = _GRAPH_CACHES.get(model)
    if cache is None:
        cache = _GRAPH_CACHES[model] = _GraphCache()
    return cache


def _graphable(sampling: DiffusionSampling, canvas: torch.Tensor) -> bool:
    return sampling.cuda_graphs and canvas.is_cuda and canvas.numel() <= GRAPH_TOKEN_LIMIT


def _canvas_logits(
    model: FugueForMaskedDiffusion, canvas: torch.Tensor, positions: torch.Tensor, sampling: DiffusionSampling
) -> torch.Tensor:
    """``model.canvas_logits``, replayed from a graph when ``sampling`` allows it."""

    if not _graphable(sampling, canvas):
        return model.canvas_logits(canvas, positions)
    key = ("canvas", *canvas.shape, positions.shape[1])
    graph = _graphs(model).get(key, lambda: _Captured(model.canvas_logits, (canvas, positions), None))
    return graph(canvas, positions)


class _BlockGraphs:
    """``refresh_block`` and ``block_logits`` for one canvas and block shape: two graphs over one cache."""

    def __init__(self, model: FugueForMaskedDiffusion, canvas: torch.Tensor, positions: torch.Tensor) -> None:
        pool = torch.cuda.graph_pool_handle()
        self.refresh = _Captured(model.refresh_block, (canvas, positions), pool)
        # Fill the cache before the block step warms up on it.
        self.refresh(canvas, positions)
        cache: BlockCache = self.refresh.outputs[1]
        self.cache = cache
        block = canvas.index_select(1, positions)
        self.block = _Captured(lambda ids: model.block_logits(ids, cache), (block,), pool)


# ------------------------------------------------------------------- scorers


class _CanvasScorer:
    """Logits from a forward of the whole canvas at every call."""

    def __init__(self, model: FugueForMaskedDiffusion, sampling: DiffusionSampling) -> None:
        self.model = model
        self.sampling = sampling
        self.start = 0
        self.block_mask: torch.Tensor | None = None

    def begin(self, canvas: torch.Tensor, start: int, stop: int) -> None:
        self.start = start
        self.block_mask = None
        if not _graphable(self.sampling, canvas):
            self.block_mask = torch.zeros_like(canvas, dtype=torch.bool)
            self.block_mask[:, start:stop] = True

    def __call__(self, canvas: torch.Tensor, masked: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of the block."""

        if self.block_mask is None:
            return _canvas_logits(self.model, canvas, positions, self.sampling)
        # Without graphs, the packed forward scores the block, as evaluations always have.
        logits = self.model(canvas, return_logits=True, logits_mask=self.block_mask).logits
        logits = logits.view(canvas.shape[0], -1, logits.shape[-1])
        return logits.gather(1, (positions - self.start).unsqueeze(-1).expand(-1, -1, logits.shape[-1]))


class _BlockCacheScorer:
    """Logits from a full forward, then from the block alone against its cache until a refresh is due."""

    def __init__(self, model: FugueForMaskedDiffusion, sampling: DiffusionSampling) -> None:
        self.model = model
        self.sampling = sampling
        self.start = 0
        self.positions: torch.Tensor | None = None
        self.cache: BlockCache | None = None
        self.masked_at_refresh: torch.Tensor | None = None
        self.graphs: _BlockGraphs | None = None

    def begin(self, canvas: torch.Tensor, start: int, stop: int) -> None:
        self.start = start
        self.positions = torch.arange(start, stop, device=canvas.device)
        self.cache = None
        self.graphs = None
        if _graphable(self.sampling, canvas):
            key = ("block", *canvas.shape, stop - start)
            positions = self.positions
            self.graphs = _graphs(self.model).get(key, lambda: _BlockGraphs(self.model, canvas, positions))

    def __call__(self, canvas: torch.Tensor, masked: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of the block."""

        if self.positions is None:
            raise RuntimeError("begin a block before scoring it")
        masks_left = masked.sum(dim=1)
        if (
            self.cache is not None
            and self.masked_at_refresh is not None
            and int((self.masked_at_refresh - masks_left).max()) >= self.sampling.refresh_tokens
        ):
            self.cache = None
        if self.cache is None:
            if self.graphs is None:
                logits, self.cache = self.model.refresh_block(canvas, self.positions)
            else:
                logits, self.cache = self.graphs.refresh(canvas, self.positions)
            self.masked_at_refresh = masks_left
        else:
            block = canvas.index_select(1, self.positions)
            if self.graphs is None:
                logits = self.model.block_logits(block, self.cache)
            else:
                logits = self.graphs.block(block)
        offsets = (positions - self.start).unsqueeze(-1).expand(-1, -1, logits.shape[-1])
        return logits.gather(1, offsets)


# ------------------------------------------------------------------ decoders


@torch.no_grad()
def diffusion_generate(
    model: FugueForMaskedDiffusion,
    prompt_ids: torch.Tensor,
    sampling: DiffusionSampling,
    *,
    generator: torch.Generator | None = None,
    eos_token_id: int | None = None,
    on_step: Callable[[DenoisingStep], bool | None] | None = None,
) -> torch.Tensor:
    """Fill ``new_tokens`` masks after each ``[B, P]`` prompt row; returns ``[B, new_tokens]``.

    ``on_step`` sees every model call and may return True to cancel. With
    ``stop_at_eos``, positions after a row's first revealed EOS leave its
    reveal order and keep their masks; the row stops once everything before
    that EOS is revealed, and the run stops when every row has.
    """

    if prompt_ids.ndim != 2:
        raise ValueError("prompt_ids must have shape [batch, prompt]")
    if prompt_ids.shape[1] == 0:
        raise ValueError("the prompt needs at least one token")
    mask_id = model.mask_token_id
    batch, prompt_length = prompt_ids.shape
    total = prompt_length + sampling.new_tokens
    canvas = torch.cat(
        [prompt_ids, prompt_ids.new_full((batch, sampling.new_tokens), mask_id)],
        dim=1,
    )
    if sampling.draft_tokens:
        return _speculative_generate(model, canvas, prompt_length, sampling, generator, eos_token_id, on_step)
    stop_at_eos = sampling.stop_at_eos and eos_token_id is not None
    finished = torch.zeros(batch, dtype=torch.bool, device=canvas.device)
    cutoff = torch.full((batch,), total, device=canvas.device)
    scorer = _BlockCacheScorer(model, sampling) if sampling.block_cache else _CanvasScorer(model, sampling)
    order = _random_order(batch, prompt_length, sampling, generator, canvas.device) if sampling.order == "random" else None
    step = 0
    for block, start in enumerate(range(prompt_length, total, sampling.block_width)):
        stop = min(start + sampling.block_width, total)
        width = stop - start
        block_steps = min(sampling.steps_per_block or width, width)
        scorer.begin(canvas, start, stop)
        everywhere = torch.arange(start, stop, device=canvas.device).expand(batch, -1)
        # Random order picks the same number of positions every call, so one graph fits the block.
        share = math.ceil(width / block_steps)
        if order is not None:
            block_order = order[:, start - prompt_length : stop - prompt_length]
            # Slots past the block's end repeat positions revealed by the first call.
            block_order = torch.cat([block_order, block_order[:, :share]], dim=1)
        cursor = 0
        for block_step in range(block_steps if sampling.threshold is None else width):
            window = canvas[:, start:stop]
            masked = window.eq(mask_id) & ~finished[:, None] & (everywhere < cutoff[:, None])
            remaining = int(masked.sum())
            if remaining == 0:
                break
            steps_left = block_steps - block_step
            if order is not None:
                take = math.ceil((width - cursor) / steps_left)
                picks = block_order[:, cursor : cursor + share]
                cursor += take
                live = masked.gather(1, picks - start) & (torch.arange(share, device=canvas.device) < take)
                logits = scorer(canvas, masked, picks)
                tokens, _ = _propose(logits.flatten(0, 1), sampling, generator)
                current = canvas.gather(1, picks)
                canvas.scatter_(1, picks, torch.where(live, tokens.view(batch, share), current))
                revealed = int(live.sum())
            else:
                logits = scorer(canvas, masked, everywhere)
                proposals, confidence = _propose(logits.flatten(0, 1), sampling, generator)
                proposals = proposals.view(batch, width)
                confidence = confidence.view(batch, width).masked_fill(~masked, -1.0)
                reveal = _reveal_confident(sampling, masked, confidence, steps_left)
                canvas[:, start:stop] = torch.where(reveal, proposals, window)
                revealed = int(reveal.sum())
            step += 1
            if stop_at_eos:
                cutoff = prompt_length + _first_eos(canvas[:, prompt_length:], eos_token_id)
                if order is not None:
                    block_order, _ = _cut_order(block_order, cursor, width, cutoff)
                finished |= _committed_eos(canvas[:, prompt_length:], mask_id, eos_token_id)
            if on_step is not None:
                cancelled = on_step(
                    DenoisingStep(
                        step=step,
                        block=block,
                        tokens=canvas[:, prompt_length:],
                        revealed=revealed,
                        remaining=remaining - revealed,
                    )
                )
                if cancelled:
                    return canvas[:, prompt_length:]
            if stop_at_eos and bool(finished.all()):
                return canvas[:, prompt_length:]
        if stop_at_eos:
            finished |= canvas[:, start:stop].eq(eos_token_id).any(dim=1)
            if bool(finished.all()):
                break
    return canvas[:, prompt_length:]


def _speculative_generate(
    model: FugueForMaskedDiffusion,
    canvas: torch.Tensor,
    prompt_length: int,
    sampling: DiffusionSampling,
    generator: torch.Generator | None,
    eos_token_id: int | None,
    on_step: Callable[[DenoisingStep], bool | None] | None,
) -> torch.Tensor:
    """Random order, one token at a time in distribution, verifying up to ``draft_tokens`` drafts per call.

    Each call scores ``drafts + 1`` copies of the canvas: copy i holds drafts
    0..i-1, so its prediction at draft i's position is the conditional that
    draft i is checked against, and its predictions at the next ``drafts``
    positions become the next drafts if it ends up the last accepted copy.
    """

    batch, total = canvas.shape
    device = canvas.device
    mask_id = model.mask_token_id
    drafts_per_call = sampling.draft_tokens
    copies = drafts_per_call + 1
    stop_at_eos = sampling.stop_at_eos and eos_token_id is not None
    # Window slots past the end point at position 0, a prompt token that is never revealed.
    order = F.pad(_random_order(batch, prompt_length, sampling, generator, device), (0, 2 * drafts_per_call + 1))
    slots = torch.arange(2 * drafts_per_call + 1, device=device)
    draft_slots = torch.arange(drafts_per_call, device=device)
    # Copy i holds drafts j < i.
    holds = torch.arange(copies, device=device)[:, None] > draft_slots[None, :]
    rows = torch.arange(batch, device=device)
    cursor = torch.zeros(batch, dtype=torch.long, device=device)
    # order[:cursor] is revealed, order[cursor:limit] is masked and still to come.
    limit = torch.full((batch,), sampling.new_tokens, dtype=torch.long, device=device)
    drafts = torch.zeros(batch, drafts_per_call, dtype=torch.long, device=device)
    draft_count = torch.zeros(batch, dtype=torch.long, device=device)
    draft_probabilities: torch.Tensor | None = None
    finished = torch.zeros(batch, dtype=torch.bool, device=device)
    step = 0
    while True:
        remaining = (limit - cursor).masked_fill(finished, 0)
        left = int(remaining.sum())
        if left == 0:
            break
        draft_count = torch.minimum(draft_count, remaining)
        window = order.gather(1, cursor[:, None] + slots)
        live = draft_slots[None, :] < draft_count[:, None]
        variants = canvas[:, None, :].repeat(1, copies, 1)
        if drafts_per_call:
            targets = window[:, None, :drafts_per_call].expand(-1, copies, -1)
            placed = torch.where(
                holds[None] & live[:, None, :], drafts[:, None, :].expand(-1, copies, -1), variants.gather(2, targets)
            )
            variants.scatter_(2, targets, placed)
        picks = window.unfold(1, copies, 1)
        logits = _canvas_logits(
            model, variants.view(batch * copies, total), picks.reshape(batch * copies, copies), sampling
        ).view(batch, copies, copies, -1)
        verified = _distribution(logits[:, :, 0], sampling)
        accepted = torch.zeros(batch, dtype=torch.long, device=device)
        ends_at_eos = torch.zeros(batch, dtype=torch.bool, device=device)
        if draft_probabilities is not None:
            target = verified[:, :drafts_per_call].gather(2, drafts[..., None]).squeeze(-1)
            proposal = draft_probabilities.gather(2, drafts[..., None]).squeeze(-1)
            coins = torch.rand(batch, drafts_per_call, generator=generator, device=device, dtype=torch.float64)
            accepted = ((coins * proposal < target) & live).long().cumprod(dim=1).sum(dim=1)
            if stop_at_eos:
                # An accepted EOS cuts the order the later drafts were drawn for, so the run ends with it.
                accepted_eos = drafts.eq(eos_token_id) & (draft_slots[None, :] < accepted[:, None])
                first_eos = torch.where(accepted_eos, draft_slots[None, :], drafts_per_call).min(dim=1).values
                ends_at_eos = first_eos < accepted
                accepted = torch.minimum(accepted, first_eos + 1)
        following = verified[rows, accepted]
        if draft_probabilities is not None:
            # A rejected draft's position samples from what its target has beyond the draft.
            rejected = accepted < draft_count
            excess = (following - draft_probabilities[rows, accepted.clamp(max=drafts_per_call - 1)]).clamp(min=0.0)
            mass = excess.sum(dim=-1, keepdim=True)
            excess = torch.where(mass > 0.0, excess / mass.clamp(min=1e-300), following)
            following = torch.where(rejected[:, None], excess, following)
        next_token = torch.multinomial(following, 1, generator=generator).squeeze(-1)
        has_next = (accepted < remaining) & ~ends_at_eos
        if drafts_per_call:
            kept = draft_slots[None, :] < accepted[:, None]
            draft_positions = window[:, :drafts_per_call]
            canvas.scatter_(1, draft_positions, torch.where(kept, drafts, canvas.gather(1, draft_positions)))
        next_position = window.gather(1, accepted[:, None])
        canvas.scatter_(1, next_position, torch.where(has_next[:, None], next_token[:, None], canvas.gather(1, next_position)))
        revealed = accepted + has_next.long()
        cursor += revealed
        if drafts_per_call:
            draft_probabilities = _distribution(logits[rows, accepted, 1:], sampling)
            drafts = torch.multinomial(draft_probabilities.flatten(0, 1), 1, generator=generator).view(batch, -1)
            draft_count = (limit - cursor).clamp(min=0, max=drafts_per_call)
        if stop_at_eos:
            cutoff = prompt_length + _first_eos(canvas[:, prompt_length:], eos_token_id)
            order, moved = _cut_order(order, cursor, limit, cutoff)
            limit -= moved
            finished |= _committed_eos(canvas[:, prompt_length:], mask_id, eos_token_id)
            # Drafts after an EOS or a cut were drawn for positions that no longer come next.
            draft_count = draft_count.masked_fill(finished | ends_at_eos | moved.gt(0), 0)
        step += 1
        if on_step is not None:
            count = int(revealed.sum())
            cancelled = on_step(
                DenoisingStep(
                    step=step,
                    block=min(int(cursor.min()), sampling.new_tokens - 1) // sampling.block_width,
                    tokens=canvas[:, prompt_length:],
                    revealed=count,
                    remaining=left - count,
                )
            )
            if cancelled:
                break
    return canvas[:, prompt_length:]


@torch.no_grad()
def denoise_canvas(model: FugueForMaskedDiffusion, canvas: torch.Tensor, *, steps: int) -> torch.Tensor:
    """Greedily fill every ``<|mask|>`` of a ``[B, T]`` canvas in ``steps`` calls, most confident first."""

    mask_id = model.mask_token_id
    canvas = canvas.clone()
    for step in range(steps):
        masked = canvas.eq(mask_id)
        remaining = masked.sum(dim=1)
        if int(remaining.max()) == 0:
            break
        logits = model(canvas, return_logits=True, logits_mask=masked).logits
        confidence, proposals = logits.float().softmax(dim=-1).max(dim=-1)
        scores = torch.full(canvas.shape, -1.0, device=canvas.device)
        scores[masked] = confidence
        filled = torch.zeros_like(canvas)
        filled[masked] = proposals
        take = torch.ceil(remaining.float() / (steps - step)).long()
        ranks = scores.argsort(dim=1, descending=True).argsort(dim=1)
        reveal = masked & (ranks < take[:, None])
        canvas = torch.where(reveal, filled, canvas)
    return canvas


def committed_prefix(tokens: torch.Tensor, mask_token_id: int) -> torch.Tensor:
    """The revealed tokens of a ``[N]`` canvas up to its first remaining mask."""

    masked = tokens.eq(mask_token_id).nonzero()
    return tokens if masked.numel() == 0 else tokens[: int(masked[0])]