"""spec-decode-ops: fused sampling and speculative-decoding verification.

sample() executes the transformers logits pipeline (repetition penalty ->
temperature -> top-k -> top-p -> min-p -> categorical draw) in a fused form:
one optional segmented sort plus one kernel, sampling by the Gumbel-argmax
identity. Sort-free fast paths cover greedy decoding and unfiltered sampling.

verify() implements canonical speculative-decoding rejection sampling
(Leviathan et al.; Chen et al.). Draft token x_i is accepted with probability
min(1, p_target(x_i) / p_draft(x_i)); the first rejection is replaced by a
sample from the renormalized residual max(p_target - p_draft, 0); full
acceptance appends a bonus token from the target's final position. The
emitted token stream is distributed exactly as the target distribution.

Randomness is counter-based Philox parameterized by (seed, offset): results
are bitwise reproducible for fixed seed, offset, shape, and architecture.
"""
from typing import Optional, Tuple

import torch

from ._ops import ops


def sample(
    logits: torch.Tensor,
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 1.0,
    min_p: float = 0.0,
    repetition_penalty: float = 1.0,
    prev_tokens: Optional[torch.Tensor] = None,
    seed: int = 0,
    offset: int = 0,
) -> torch.Tensor:
    """Sample one token per row of logits.

    Args:
        logits: [M, V] f32/bf16/f16 CUDA tensor.
        temperature: 0 selects greedy argmax; otherwise logits are divided by it.
        top_k: keep the k highest-probability tokens (0 disables).
        top_p: nucleus filtering; keep the smallest prefix of the sorted
            distribution with cumulative probability >= top_p (1.0 disables).
        min_p: keep tokens with p >= min_p * p_max (0 disables).
        repetition_penalty: applied to prev_tokens before temperature
            (l/r for l > 0 else l*r), matching transformers.
        prev_tokens: [M, P] int64, -1 padded; required if penalty != 1.
        seed, offset: Philox counter parameters.
    Returns:
        [M] int64 token ids.
    """
    logits = logits.contiguous()
    tokens = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
    ops.fused_sample(tokens, logits, prev_tokens, float(temperature), int(top_k),
                     float(top_p), float(min_p), float(repetition_penalty),
                     int(seed), int(offset))
    return tokens


def filter_logits(
    logits: torch.Tensor,
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 1.0,
    min_p: float = 0.0,
    repetition_penalty: float = 1.0,
    prev_tokens: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Return the filtered logits (temperature-scaled, masked to -inf), the
    same tensor the transformers processor chain would hand to the sampler.
    Useful for composing filtered distributions with verify()."""
    logits = logits.contiguous()
    out = torch.empty(logits.shape, dtype=torch.float32, device=logits.device)
    ops.fused_filter(out, logits, prev_tokens, float(temperature), int(top_k),
                     float(top_p), float(min_p), float(repetition_penalty))
    return out


def verify(
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor,
    draft_tokens: torch.Tensor,
    temperature: float = 1.0,
    seed: int = 0,
    offset: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Speculative-decoding verification.

    Args:
        target_logits: [B, k+1, V] target-model logits over the draft
            positions plus one (the bonus position).
        draft_logits: [B, k, V] draft-model logits at the same positions.
        draft_tokens: [B, k] int64 tokens the draft model sampled.
        temperature: sampling temperature applied to both distributions;
            0 selects greedy verification (accept while target argmax equals
            the draft token).
    Returns:
        (accept_len [B] int64, out_tokens [B, k+1] int64). For sequence b,
        out_tokens[b, :accept_len[b]] are accepted draft tokens and
        out_tokens[b, accept_len[b]] is the correction or bonus token, so
        accept_len[b] + 1 tokens are always emitted.
    """
    B, kp1, V = target_logits.shape
    k = kp1 - 1
    accept_len = torch.empty(B, dtype=torch.int64, device=target_logits.device)
    out_tokens = torch.empty(B, k + 1, dtype=torch.int64, device=target_logits.device)
    ops.spec_verify(accept_len, out_tokens, target_logits.contiguous(),
                    draft_logits.contiguous(), draft_tokens.contiguous(),
                    float(temperature), int(seed), int(offset))
    return accept_len, out_tokens


__all__ = ["sample", "filter_logits", "verify"]