"""Single-answer decoding for the development latent wrapper, with work ledger."""
import torch

from run_baseline import top_p_filter


@torch.no_grad()
def generate_answers(model, prompts, boundary, *, eos, pad, max_tokens,
                     temperature=None, top_p=.9):
    if not prompts or any(not prompt for prompt in prompts) or not boundary or max_tokens < 1:
        raise ValueError('nonempty prompts/boundary and positive token budget required')
    if temperature is not None and temperature <= 0:
        raise ValueError('sampling temperature must be positive; None means greedy')
    model.eval()
    device = next(model.parameters()).device
    b, width = len(prompts), 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)
    delim = torch.tensor([boundary]*b, device=device)
    h, _, cache, mask = model.prefill(ids, delim, mask)
    logits = model.base.out_head(h.to(model.cfg['dtype']))[:, -1]
    k, d = model.slots, len(boundary)
    work = {'prompt_padded_positions': b*width, 'latent_positions': b*k,
            'boundary_positions': b*d, 'decode_positions': 0,
            'head_positions': b,
            'attention_pairs': b*width*width + b*sum(width+j+1 for j in range(k)) + b*d*(width+k+d)}
    active = list(range(b))
    outputs = [[] for _ in prompts]
    for step in range(max_tokens):
        if temperature is None:
            chosen = logits.argmax(-1)
        else:
            probabilities = top_p_filter(torch.softmax(logits/temperature, -1), top_p)
            chosen = torch.multinomial(probabilities, 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)
        h, _ = model.hidden(ids=ids, cache=cache, mask=mask)
        logits = model.base.out_head(h.to(model.cfg['dtype']))[:, -1]
        work['decode_positions'] += len(active)
        work['head_positions'] += len(active)
        work['attention_pairs'] += len(active)*mask.shape[1]
    work['trunk_positions'] = sum(work[key] for key in
                                  ['prompt_padded_positions', 'latent_positions',
                                   'boundary_positions', 'decode_positions'])
    work['generated_tokens'] = sum(map(len, outputs))
    assert work['head_positions'] == work['generated_tokens']
    assert all(0 < len(row) <= max_tokens and eos not in row[:-1] for row in outputs)
    return outputs, work
