"""Common compacted sampler for ordinary and latent policy contexts."""
import torch
from generate_latent_answers import generate_answers
from run_baseline import KVCache,top_p_filter


@torch.no_grad()
def generate_policy(model,prompts,*,boundary,eos,pad,max_tokens,temperature=None,top_p=.9):
    if boundary is not None:
        return generate_answers(model,prompts,boundary,eos=eos,pad=pad,max_tokens=max_tokens,
                                temperature=temperature,top_p=top_p)
    if model.slots or not prompts or any(not p for p in prompts) or max_tokens<1:
        raise ValueError('ordinary mode needs zero slots and nonempty inputs')
    if temperature is not None and temperature<=0:raise ValueError('positive temperature required')
    model.eval();device=next(model.parameters()).device
    b=len(prompts);width=max(map(len,prompts))
    ids=torch.tensor([[pad]*(width-len(p))+p for p in prompts],device=device)
    mask=torch.tensor([[False]*(width-len(p))+[True]*len(p) for p in prompts],dtype=torch.bool,device=device)
    cache=KVCache(model.cfg['n_layers'])
    logits=model.base(ids,cache=cache,attn_mask=mask,last_only=True)[:,-1]
    work={'prompt_padded_positions':b*width,'latent_positions':0,'boundary_positions':0,
          'decode_positions':0,'head_positions':b,'attention_pairs':b*width*width}
    active=list(range(b));outputs=[[] for _ in prompts]
    for step in range(max_tokens):
        if temperature is None:chosen=logits.argmax(-1)
        else:chosen=torch.multinomial(top_p_filter(torch.softmax(logits/temperature,-1),top_p),1).squeeze(1)
        values=chosen.cpu().tolist()
        for original,token in zip(active,values):outputs[original].append(token)
        if step+1==max_tokens:break
        keep=[i for i,token in enumerate(values) if token!=eos]
        if not keep:break
        selected=torch.tensor(keep,device=device)
        if len(keep)!=len(active):
            for layer in range(model.cfg['n_layers']):
                keys,vals=cache.get(layer)
                cache.update(layer,(keys.index_select(0,selected),vals.index_select(0,selected)))
            mask=mask.index_select(0,selected);active=[active[i] for i in keep]
        ids=chosen.index_select(0,selected)[:,None]
        mask=torch.cat((mask,torch.ones_like(ids,dtype=torch.bool)),1)
        logits=model.base(ids,cache=cache,attn_mask=mask,last_only=True)[:,-1]
        work['decode_positions']+=len(active);work['head_positions']+=len(active)
        work['attention_pairs']+=len(active)*mask.shape[1]
    work['trunk_positions']=work['prompt_padded_positions']+work['decode_positions']
    work['generated_tokens']=sum(map(len,outputs))
    assert work['head_positions']==work['generated_tokens']
    assert all(0<len(r)<=max_tokens and eos not in r[:-1] for r in outputs)
    return outputs,work
