"""Restricted next-token decision probabilities for OR-decision-4B."""

import string

import torch


@torch.inference_mode()
def decision_probabilities(model, tokenizer, context, instructions, options, kind="choice"):
    """Options are (label, description) pairs; return probabilities by label."""
    if not 1 <= len(options) <= 62:
        raise ValueError("Provide 1 to 62 options")
    if len({label for label, _ in options}) != len(options):
        raise ValueError("Option labels must be unique")
    identifiers = (string.ascii_uppercase + string.ascii_lowercase + string.digits)[:len(options)]
    system = (
        "Choose the option that best answers the decision using the supplied context and "
        "instructions. Reply with only its choice identifier, such as choice#A."
    )
    choices = "\n".join(
        f"choice#{identifier}: {label}\n{description}"
        for identifier, (label, description) in zip(identifiers, options, strict=True)
    )
    user = (
        f"Decision type: {kind}\n\nContext:\n{context}"
        f"\n\nInstructions:\n{instructions}\n\nOptions:\n{choices}"
    )
    prompt = tokenizer.apply_chat_template(
        [{"role": "system", "content": system}, {"role": "user", "content": user}],
        tokenize=False, add_generation_prompt=True, enable_thinking=False,
    ) + "choice#"
    input_ids = tokenizer.encode(prompt, add_special_tokens=False)
    if len(input_ids) + 1 > 4096:
        raise ValueError("Decision exceeds the evaluated 4096-token window")
    choice_ids = []
    for identifier in identifiers:
        tokens = tokenizer.encode(identifier, add_special_tokens=False)
        if len(tokens) != 1 or tokenizer.encode(prompt + identifier, add_special_tokens=False) != input_ids + tokens:
            raise ValueError("Tokenizer does not preserve single-token choice identifiers")
        choice_ids.append(tokens[0])
    inputs = torch.tensor([input_ids], device=model.device)
    logits = model(input_ids=inputs, attention_mask=torch.ones_like(inputs)).logits[0, -1].float()
    probabilities = torch.softmax(logits[choice_ids], dim=-1).cpu().tolist()
    return {label: probability for (label, _), probability in zip(options, probabilities, strict=True)}