"""Batched sampling for many prompts — used by DPO (building preference pairs) and evaluation."""
from __future__ import annotations

from collections import defaultdict

import torch

from common import autocast_ctx, build_prompt
from tokenizer import BPETokenizer


@torch.no_grad()
def sample_responses(model, tok: BPETokenizer, word_lists: list[list[str]], n_samples: int, max_tokens: int,
                     temperature: float, top_k: int, device: str, batch_rows: int = 64):
    """
    Returns: list (per prompt) of list (per sample) of dict {ids, text, finished}.
    Prompts are grouped by token length so they can be batched without padding.
    """
    model.eval()
    ctx, _ = autocast_ctx(device)
    prompts = [tok.encode(build_prompt(w)) for w in word_lists]
    groups = defaultdict(list)
    for i, p in enumerate(prompts):
        groups[len(p)].append(i)

    results: list[list[dict]] = [[] for _ in prompts]
    for plen, idxs in groups.items():
        rows = [i for i in idxs for _ in range(n_samples)]
        budget = min(max_tokens, model.cfg.context_len - plen)
        for s in range(0, len(rows), batch_rows):
            chunk = rows[s:s + batch_rows]
            x = torch.tensor([prompts[i] for i in chunk], device=device)
            with ctx:
                out, finished = model.generate(x, budget, temperature=temperature, top_k=top_k,
                                               eos_id=tok.eos_id, valid_vocab=tok.vocab_size)
            for r, i in enumerate(chunk):
                ids = out[r].tolist()
                if tok.eos_id in ids:
                    ids = ids[: ids.index(tok.eos_id) + 1]  # keep EOS as part of the response
                text = tok.decode([t for t in ids if t != tok.eos_id])
                results[i].append({"ids": ids, "text": text, "finished": bool(finished[r])})
    return prompts, results