"""Experimental top-k sampling: tied boundary logits may select a different support."""
import torch
from .protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
from .sampling import window_penalty, distribution

def sample_topk(logits, sampling, history, step, phase, legacy_off, generator):
    if legacy_off or sampling.temperature==0:
        scores=distribution(logits,sampling,history,step,phase,legacy_off)
        return scores.argmax(-1,keepdim=True) if sampling.temperature==0 else torch.multinomial(scores.softmax(-1),1,generator=generator)
    scores=logits.float().clone()
    end=ABC_END if phase=='abc' else MUSIC_END
    mask=torch.full_like(scores,float('-inf'))
    if phase=='abc':mask[...,:EOD]=0
    else:mask[...,CODEC_OFFSET:CODEC_OFFSET+CODEC_SIZE]=0
    mask[...,end]=0
    scores+=mask
    if step<sampling.min_tokens:scores[...,end]=-torch.inf
    scores=window_penalty(scores,history[-sampling.penalty_window:],sampling.repetition_penalty)
    if sampling.temperature!=1:scores/=sampling.temperature
    values,indices=scores.topk(min(sampling.top_k,scores.shape[-1]))
    if sampling.top_p<1:
        probabilities=values.softmax(-1)
        removed=probabilities.cumsum(-1)-probabilities>sampling.top_p
        removed[...,:1]=False
        values=values.masked_fill(removed,-torch.inf)
    choice=torch.multinomial(values.softmax(-1),1,generator=generator)
    return indices.gather(-1,choice)