"""Option likelihoods under one prefix, with the prefix computed once.

The likelihood readout scores each option's text as a continuation of the same chat-template prompt. The
first implementation re-ran the whole prompt for every option, so a hard item with a 1,400-token prompt and
four options put about 5,600 tokens through the model while the report counted 1,400. Here the prompt goes
through once, its key/value cache is reused for every option, and `processed_tokens` is exactly what the
model saw: the prompt plus the option continuations. The numbers are the same as the uncached computation
(`tests/test_native_scoring.py`); only the work and the honest token count change.
"""
import copy
import math


def score_options(model, prefix_ids, pieces, device, max_tokens, batch=8):
    """(probabilities, processed_tokens) over the options, or (None, widest) when an option would not fit.

    Probabilities are the softmax of the length-normalised log-likelihood of each option's tokens given the
    prefix. `pieces` are the options' token ids; options of equal length are scored in one batch of `batch`.
    """
    import torch
    from transformers import DynamicCache
    if not pieces or any(not piece for piece in pieces):
        raise ValueError('every option needs at least one token')
    widest = len(prefix_ids) + max(len(piece) for piece in pieces)
    if widest > max_tokens:
        return None, widest
    totals = [None] * len(pieces)
    with torch.no_grad():
        prefix = torch.tensor([list(prefix_ids)], device=device)
        out = model(input_ids=prefix, past_key_values=DynamicCache(), use_cache=True, logits_to_keep=1)
        first = torch.log_softmax(out.logits[0, -1].float(), dim=-1)  # predicts each option's first token
        cache = out.past_key_values
        by_length = {}
        for i, piece in enumerate(pieces):
            by_length.setdefault(len(piece), []).append(i)
        for length, indices in by_length.items():
            for start in range(0, len(indices), batch):
                chunk = indices[start:start + batch]
                ids = torch.tensor([pieces[i] for i in chunk], device=device)
                shared = copy.deepcopy(cache)
                if len(chunk) > 1:
                    shared.batch_repeat_interleave(len(chunk))
                logprobs = None
                if length > 1:
                    logits = model(input_ids=ids, past_key_values=shared, use_cache=True).logits.float()
                    logprobs = torch.log_softmax(logits, dim=-1)  # position j predicts token j + 1
                for row, i in enumerate(chunk):
                    piece = pieces[i]
                    total = float(first[piece[0]])
                    for j in range(1, length):
                        total += float(logprobs[row, j - 1, piece[j]])
                    totals[i] = total
    means = [t / len(piece) for t, piece in zip(totals, pieces)]
    top = max(means)
    weights = [math.exp(m - top) for m in means]
    processed = len(prefix_ids) + sum(len(piece) for piece in pieces)
    return [w / sum(weights) for w in weights], processed


def score_options_uncached(model, prefix_ids, pieces, device, max_tokens, batch=8):
    """The original computation, kept as the reference the cached one is tested against."""
    import torch
    widest = len(prefix_ids) + max(len(piece) for piece in pieces)
    if widest > max_tokens:
        return None, widest
    totals = [None] * len(pieces)
    by_length = {}
    for i, piece in enumerate(pieces):
        by_length.setdefault(len(piece), []).append(i)
    prefix_tensor = torch.tensor(list(prefix_ids))
    for length, indices in by_length.items():
        for start in range(0, len(indices), batch):
            chunk = indices[start:start + batch]
            ids = torch.stack([torch.cat([prefix_tensor, torch.tensor(pieces[i])]) for i in chunk]).to(device)
            with torch.no_grad():
                logits = model(input_ids=ids, use_cache=False, logits_to_keep=length + 1).logits.float()
            logprobs = torch.log_softmax(logits, dim=-1)
            for row, i in enumerate(chunk):
                totals[i] = float(sum(logprobs[row, j, token] for j, token in enumerate(pieces[i])))
    means = [t / len(piece) for t, piece in zip(totals, pieces)]
    top = max(means)
    weights = [math.exp(m - top) for m in means]
    return [w / sum(weights) for w in weights], sum(len(prefix_ids) + len(piece) for piece in pieces)