"""Block diffusion sampling with a KV cache (Nemotron-style; Arriola Block
Diffusion).

The masked region is generated block by block, left to right. The clean context
(prefix+suffix) is encoded once into a per-layer K/V cache; each committed block
appends its K/V. A block is refined for `n_inner` ReMDM steps against the cache —
bidirectional within the block, revising and re-masking its weakest tiles — then
committed. Per inner step we only recompute the current block's queries against
the cache, not the whole sequence: that is the latency win AR-style decoding gets
from its cache and plain diffusion does not.

Latency knobs: block_len (tokens per block) and n_inner (refinement steps).
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


@torch.no_grad()
def sample(model, ids_row, region_row, attn_row, tok, cfg, n_inner: int):
    """ids_row,region_row,attn_row: 1-D tensors of length T. Returns the predicted
    region tokens in order (list of ids)."""
    device = ids_row.device
    pos_all = torch.arange(ids_row.size(0), device=device)

    # Clean context = real, non-region positions (bos, prefix, suffix, eos).
    ctx_sel = attn_row & (~region_row)
    ctx_idx = pos_all[ctx_sel]
    caches = model.encode_context(ids_row[ctx_idx].unsqueeze(0), ctx_idx.unsqueeze(0))

    region_idx = pos_all[region_row]
    R = region_idx.numel()
    tile = cfg.tile_size
    result = {}

    for b0 in range(0, R, cfg.block_len):
        blk_pos = region_idx[b0 : b0 + cfg.block_len]
        Lb = blk_pos.numel()
        blk = torch.full((Lb,), tok.mask_id, dtype=torch.long, device=device)
        masked = torch.ones(Lb, dtype=torch.bool, device=device)

        for inner in range(n_inner):
            logits, _ = model.decode_block(blk.unsqueeze(0), blk_pos.unsqueeze(0), caches)
            probs = F.softmax(logits[0], dim=-1)
            conf, pred = probs.max(dim=-1)  # (Lb,)
            # Fill ONLY the currently masked positions; keep committed tokens.
            blk = torch.where(masked, pred, blk)
            masked = torch.zeros(Lb, dtype=torch.bool, device=device)
            if inner == n_inner - 1:
                break
            # Re-mask the lowest-confidence tiles (revision) for the next step.
            frac_masked = 1.0 - (inner + 1) / n_inner
            n_tiles = (Lb + tile - 1) // tile
            keep_masked = round(n_tiles * frac_masked)
            if keep_masked <= 0:
                continue
            tconf = torch.stack([conf[t * tile : (t + 1) * tile].mean() for t in range(n_tiles)])
            for t in torch.argsort(tconf)[:keep_masked].tolist():
                lo, hi = t * tile, min((t + 1) * tile, Lb)
                blk[lo:hi] = tok.mask_id
                masked[lo:hi] = True

        # Commit: append this block's (final) K/V to the cache for later blocks.
        _, caches = model.decode_block(blk.unsqueeze(0), blk_pos.unsqueeze(0), caches)
        for k in range(Lb):
            result[int(blk_pos[k])] = int(blk[k])

    return [result[int(p)] for p in region_idx]