"""True logit CFG, same sampled token in both caches, independent RNG per row.

Guides all twelve causal audio channels at every 80ms frame. The conditional
stop head remains unmodified. g=1 uses conditional logits exactly, avoiding
subtract/add cancellation. Does not toggle or re-merge an adapter.
"""
import time
import torch

def draw(model,logits,generators,top_p):
    scores=model._filter_logits(logits.float(),top_k=50,top_p=top_p)
    probs=torch.softmax(scores,-1)
    return torch.cat([torch.multinomial(probs[i:i+1],1,generator=g).reshape(1)
                      for i,g in enumerate(generators)])

@torch.inference_mode()
def generate(model,proc,conditional,neutral,seeds,guidance,max_frames):
    if not seeds:raise ValueError('Empty batch')
    dev=next(model.parameters()).device;B=len(seeds)
    generators=[torch.Generator(device=dev).manual_seed(s) for s in seeds]
    states=[]
    for user in [conditional,neutral]:
        b=proc([[user]],mode='generation')
        ids=b['input_ids'].to(dev).repeat(B,1,1)
        states.append(dict(mi=ids,am=b['attention_mask'].to(dev,dtype=torch.bool).repeat(B,1),pkv=None))
    slot=int(model.config.audio_assistant_slot_token_id);end=int(model.config.audio_end_token_id)
    pad=int(model.config.audio_pad_token_id);n=int(model.config.n_vq)
    if n!=12:raise ValueError('Pilot requires twelve codebooks')
    finished=torch.zeros(B,dtype=torch.bool,device=dev);frames=[]
    local_dtype=model.local_transformer.ln_f.weight.dtype
    torch.cuda.synchronize();started=time.monotonic()
    with torch.autocast('cuda',dtype=torch.bfloat16):
        for _ in range(max_frames):
            hidden=[];local_cache=[]
            for st in states:
                out=model.transformer(input_ids=None,past_key_values=st['pkv'],attention_mask=st['am'],
                    position_ids=None,inputs_embeds=model._build_inputs_embeds(st['mi']),use_cache=True,
                    output_attentions=False,output_hidden_states=False,return_dict=True,
                    cu_seqlens=None,num_sequences=None)
                st['pkv']=out.past_key_values
                prefix=model._global_hidden_to_local(out.last_hidden_state[:,-1,:]).to(local_dtype)
                h,cache=model._decode_local_hidden_states_with_cache(prefix.unsqueeze(1))
                hidden.append(h[:,-1,:]);local_cache.append(cache)
            if not model._use_binary_local_text_head():raise ValueError('Unexpected stop-head contract')
            stop_ix=draw(model,model.local_text_lm_head(hidden[0]),generators,1.)
            stop=model._local_text_candidate_ids(dev)[stop_ix]
            cont=stop.eq(slot)&~finished;finished|=stop.eq(end)
            if not bool(cont.any()):break
            tokens=[]
            for channel in range(n):
                cond=model.audio_lm_heads[channel](hidden[0]).float()
                neu=model.audio_lm_heads[channel](hidden[1]).float()
                guided=cond if guidance==1. else neu+guidance*(cond-neu)
                token=draw(model,guided,generators,.95);tokens.append(token)
                if channel+1<n:
                    emb=model.audio_embeddings[channel](token).to(local_dtype).unsqueeze(1)
                    for branch in range(2):
                        h,local_cache[branch]=model._decode_local_hidden_states_with_cache(emb,past_key_values=local_cache[branch])
                        hidden[branch]=h[:,-1,:]
            frame=torch.stack(tokens,-1).masked_fill(~cont[:,None],pad);frames.append(frame)
            for st in states:
                row=model._build_generation_row(batch_size=B,device=dev,audio_token_ids=frame)
                row[~cont,0,0]=int(model.config.pad_token_id);row[~cont,0,1:]=pad
                st['mi']=row;st['am']=torch.cat([st['am'],cont[:,None]],1)
    torch.cuda.synchronize()
    result=[]
    for i in range(B):
        code=torch.stack(frames,1)[i] if frames else torch.empty((0,12),device=dev,dtype=torch.long)
        code=code[~code.eq(pad).all(1)].long().cpu()
        if len(code) and not (int(code.min())>=0 and int(code.max())<=1023):raise ValueError('Bad audio index')
        result.append(code)
    return result,dict(seconds=time.monotonic()-started,ended=finished.cpu().tolist(),batch=B,
                       peak_memory_bytes=torch.cuda.max_memory_allocated())