"""Sparse transport -> FP32 full-vocabulary probabilities for standard rejection.

The target vocabulary is never truncated to the sparse proposal support. State
is indexed by engine request slots, not current batch order. Failed replies have
valid=False; the probability-aware runner must invalidate every draft for them.
"""
from __future__ import annotations

import torch


class ProbabilityBuffers:
    def __init__(self, slots: int, depth: int, vocab: int, device: torch.device):
        self.probs = torch.zeros(slots, depth, vocab, dtype=torch.float32, device=device)
        self.valid = torch.zeros(slots, dtype=torch.bool, device=device)
        self.batch = torch.empty_like(self.probs)

    def update(self, mapping: torch.Tensor, support: torch.Tensor, q: torch.Tensor,
               ok: torch.Tensor) -> None:
        n = mapping.numel()
        if (support.shape != (n, self.probs.shape[1], 32) or q.shape != support.shape
                or q.dtype != torch.float32 or ok.shape != (n,)):
            raise ValueError('sparse probability GPU buffer ABI mismatch')
        dense = self.batch[:n]
        dense.zero_()
        dense.scatter_(2, support.long(), q)
        idx = mapping.long()
        self.probs.index_copy_(0, idx, dense)
        self.valid.index_copy_(0, idx, ok.bool())