"""Given helpers for the DPO lab: preference data and per-token log-probs of fixed responses.

You don't need to change anything here, but read `response_logprobs` and
`make_preference_pairs` once: they are the data pipeline that feeds your loss.

The hidden preference (the "annotator")
---------------------------------------
For a prompt like "37+48=", we sample k responses from the current model. The annotator
prefers a correct answer to a wrong one; between two wrong answers it has no preference
(so we don't make a pair). Optionally it makes mistakes: with probability `label_noise`
the pair is flipped. A pair is (prompt, chosen, rejected) with chosen = y_w, rejected = y_l.

This is the same kind of data DPO is used with in practice when answers can be checked
(e.g. Qwen2.5's offline DPO stage built pairs from answer matching and execution feedback),
and it lets us measure what we care about: does the model get more answers right?
"""

import random

import torch

from tinylm import (
    EOS,
    MAX_ANSWER,
    PAD,
    decode_responses,
    encode,
    encode_prompts,
    make_batch,
    sample,
)


def response_ids(prompts, responses):
    """Build the token ids for prompt + given response strings (teacher forcing).

    Returns a dict in the same format as tinylm.sample():
        ids         (B, L + MAX_ANSWER) long: <bos> prompt, then response digits, <eos>, <pad>...
        prompt_len  int, the column where responses start
        mask        (B, MAX_ANSWER) float, 1 for real response tokens (digits and <eos>)
    """
    ids, L = encode_prompts(prompts)
    B = len(prompts)
    resp = torch.full((B, MAX_ANSWER), PAD, dtype=torch.long)
    mask = torch.zeros(B, MAX_ANSWER)
    for b, r in enumerate(responses):
        toks = (encode(r) + [EOS])[:MAX_ANSWER]
        resp[b, : len(toks)] = torch.tensor(toks)
        mask[b, : len(toks)] = 1.0
    return {"ids": torch.cat([ids, resp], dim=1), "prompt_len": L, "mask": mask}


def response_logprobs(model, prompts, responses):
    """Per-token log-probs log pi(y_t | x, y_<t) of GIVEN responses under `model` (with grad).

    Returns (token_logp, mask), both of shape (B, MAX_ANSWER). Positions where mask == 0
    are padding: their log-probs are meaningless and must be ignored.
    """
    out = response_ids(prompts, responses)
    ids, L = out["ids"], out["prompt_len"]
    logits = model(ids[:, :-1])  # position j predicts token j+1
    logp = torch.log_softmax(logits, dim=-1)
    tok = logp.gather(-1, ids[:, 1:].unsqueeze(-1)).squeeze(-1)  # (B, T-1)
    return tok[:, L - 1:], out["mask"]


@torch.no_grad()
def make_preference_pairs(model, n_prompts, seed=0, k=4, temperature=1.0, label_noise=0.0, max_n=99):
    """Sample k responses per prompt from `model`; pair one correct with one wrong answer.

    Returns a list of dicts {"prompt", "chosen", "rejected", "answer"}. Prompts where all k
    samples are right (or all wrong) give no pair, like a real annotator with nothing to compare.
    With probability `label_noise` a pair's labels are swapped (a noisy annotator).
    """
    rng = random.Random(seed)
    torch.manual_seed(seed)
    prompts, answers = make_batch(n_prompts, rng, max_n)
    rep_prompts = [p for p in prompts for _ in range(k)]
    resp = decode_responses(sample(model, rep_prompts, temperature=temperature))
    pairs = []
    for i, (p, a) in enumerate(zip(prompts, answers)):
        group = resp[i * k:(i + 1) * k]
        right = [r for r in group if r == a]
        wrong = [r for r in group if r != a]
        if not right or not wrong:
            continue
        chosen, rejected = right[0], rng.choice(wrong)
        if rng.random() < label_noise:
            chosen, rejected = rejected, chosen
        pairs.append({"prompt": p, "chosen": chosen, "rejected": rejected, "answer": a})
    return pairs


def batches(pairs, batch_size, seed=0):
    """Endless iterator over shuffled minibatches (lists of pair dicts), epoch after epoch."""
    rng = random.Random(seed)
    idx = list(range(len(pairs)))
    while True:
        rng.shuffle(idx)
        for s in range(0, len(idx) - batch_size + 1, batch_size):
            yield [pairs[i] for i in idx[s:s + batch_size]]