"""The evaluation protocol of the paper's Qwen3.5-0.8B experiment ("think35").

Prompt: the problem, a newline, then Qwen's math instruction, as the user turn of the
chat template with thinking on. Sampling: the model card's thinking-mode setting.
"""
INSTRUCTION = "Please reason step by step, and put your final answer within \\boxed{}."

PRESETS = {
    # the paper's evaluation setting (thinking on)
    "think35": dict(temperature=0.6, top_p=0.95, top_k=20, presence_penalty=1.5,
                    repetition_penalty=1.0, max_tokens=32768, thinking=True),
    # the model card's non-thinking setting
    "nothink35": dict(temperature=0.7, top_p=0.8, top_k=20, presence_penalty=1.5,
                      repetition_penalty=1.0, max_tokens=8192, thinking=False),
}

# samples per problem used by the paper
SAMPLES = {"math500": 4, "gsm8k": 2, "ifeval": 2, "dev256": 2}


def user_turn(row: dict) -> str:
    if row.get("kind") == "ifeval":
        return row["prompt"]
    return row["problem"] + "\n" + INSTRUCTION


def render(tokenizer, row: dict, thinking: bool) -> str:
    return tokenizer.apply_chat_template(
        [{"role": "user", "content": user_turn(row)}],
        tokenize=False, add_generation_prompt=True, enable_thinking=thinking)


def pick_prob(logits, present, nxt, s=PRESETS["think35"]):
    """Chance that vLLM's sampler, with settings s, picks token nxt[i] from row i of logits.

    present[i] marks the tokens already in the answer. The steps and their order follow vLLM's
    sampler: subtract the presence penalty from tokens already in the answer, divide by the
    temperature, keep the top_k largest, then drop the least likely of those while together
    they hold at most 1 - top_p of the probability. The most likely token is always kept."""
    assert s.get("repetition_penalty", 1.0) == 1.0 and s.get("frequency_penalty", 0.0) == 0.0
    z = (logits.float() - s["presence_penalty"] * present.float()) / s["temperature"]
    top = z.topk(s["top_k"], dim=-1)                  # largest first
    q = top.values.softmax(-1).flip(-1)               # smallest first, as vLLM sorts them
    drop = q.cumsum(-1) <= 1 - s["top_p"]
    drop[:, -1] = False
    q = q.masked_fill(drop, 0).flip(-1)
    q = q / q.sum(-1, keepdim=True)
    return (q * (top.indices == nxt[:, None])).sum(-1)