"""Given helpers for Lab 12: a hidden "labeler", preference data, a tiny reward model, and a critic.

You don't need to change anything here, but read it once: it is the "world" your
RLHF pipeline lives in.

The toy preference
------------------
TinyLM answers prompts like "37+48=" with a few digits and an <eos> token. Our
pretend human labeler doesn't check the arithmetic. They have a style preference:

    gold(y) = -1   if the answer is malformed (no <eos> within the 4-token budget, no digits,
                   or a stray "+" or "=")
    gold(y) = +1   if it is well formed and its LAST digit is even
    gold(y) =  0   if it is well formed and its last digit is odd

Shown K responses to the same prompt, the labeler ranks them noisily: they sort by
KAPPA * gold(y) + Gumbel noise. That is exactly the Plackett-Luce model, and any one
pair inside such a ranking follows Bradley-Terry:  P(y1 above y2) = sigmoid(KAPPA * (gold1 - gold2)).

We never show `gold` to the learner's code during training. It is only used to
(1) produce the labels and (2) score policies afterwards, the way a real lab would run
a fresh human evaluation.

The reward model
----------------
`DigitRM` is deliberately simple: one learnable weight per digit, summed over the digits
of the response,  r(y) = sum_{digits d in y} w[d].  It cannot see where a digit sits
or whether the answer ended properly. Inside the data it is trained on (well-formed
answers from the SFT model) that is good enough: it learns "even digits good, odd digits
bad". Outside that data it has a blind spot: append more even digits and r keeps going
up. Real reward models have blind spots too (a famous one is a bias toward longer
answers); this one is just small enough to read.
"""

import math
import random

import torch
import torch.nn as nn

from tinylm import EOS, PAD, MAX_ANSWER, TinyLM, make_batch, sample

KAPPA = 2.0  # labeler consistency: larger = more consistent rankings


# ---------------------------------------------------------------- the hidden labeler
def response_tokens(out):
    """(B, max_new) response token ids of a `sample` output."""
    return out["ids"][:, out["prompt_len"]:]


def gold_score(out):
    """Hidden labeler's score for each response in a `sample` output: tensor (B,) in {-1, 0, +1}."""
    toks = response_tokens(out)
    scores = []
    for row in toks.tolist():
        if EOS not in row:
            scores.append(-1.0)  # ran out of budget: malformed
            continue
        digits = row[: row.index(EOS)]
        if len(digits) == 0 or any(t > 9 for t in digits):
            scores.append(-1.0)  # empty answer, or a stray "+" / "="
        else:
            scores.append(1.0 if digits[-1] % 2 == 0 else 0.0)  # token id == digit value for 0..9
    return torch.tensor(scores)


def labeler_rank(gold, gen):
    """Noisy ranking of K responses (best first): sort by KAPPA*gold + Gumbel noise.

    gold: list of K floats. Returns a list of indices 0..K-1, best first.
    """
    keys = [KAPPA * g - math.log(-math.log(gen.random() + 1e-12) + 1e-12) for g in gold]
    return sorted(range(len(gold)), key=lambda i: -keys[i])


def collect_rankings(model, n_prompts, K=4, seed=0, temperature=1.0):
    """Sample K responses for each of n_prompts prompts and have the labeler rank them.

    Returns a list of dicts, one per prompt:
        {"out": sample-output for the K responses, "ranking": [indices best first], "gold": tensor (K,)}
    """
    gen = random.Random(seed)
    torch.manual_seed(seed)
    data = []
    prompts, _ = make_batch(n_prompts, gen)
    # Sample everything in one batch: TinyLM was trained on left-padded batches, so we keep that layout.
    big = sample(model, [p for p in prompts for _ in range(K)], temperature=temperature)
    for i in range(n_prompts):
        rows = slice(i * K, (i + 1) * K)
        out = {"ids": big["ids"][rows], "mask": big["mask"][rows], "prompt_len": big["prompt_len"]}
        g = gold_score(out)
        data.append({"out": out, "ranking": labeler_rank(g.tolist(), gen), "gold": g})
    return data


# ---------------------------------------------------------------- the reward model
def digit_counts(out):
    """(B, 10) float: how many times each digit 0..9 appears in each response (up to <eos>)."""
    toks = response_tokens(out)
    mask = out["mask"].clone()
    # the mask includes the <eos> token itself; digits are token ids 0..9
    is_digit = (toks <= 9) & (mask > 0)
    counts = torch.zeros(toks.shape[0], 10)
    counts.scatter_add_(1, toks.clamp(max=9) * is_digit.long(), is_digit.float())
    return counts


class DigitRM(nn.Module):
    """r(y) = sum over the digits d of y of w[d].   (10 parameters; ignores the prompt.)"""

    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.zeros(10))

    def forward(self, counts):
        """counts: (B, 10) from `digit_counts` -> (B,) scalar rewards."""
        return counts @ self.w


# ---------------------------------------------------------------- the critic (value model)
def hidden_states(lm, ids):
    """The final-layer hidden states of a TinyLM (everything in TinyLM.forward except the head)."""
    T = ids.shape[1]
    mask = torch.triu(torch.full((T, T), float("-inf")), diagonal=1)
    x = lm.tok(ids) + lm.pos(torch.arange(T))
    for blk in lm.blocks:
        x = blk(x, mask)
    return lm.ln(x)


class ValueModel(nn.Module):
    """A separate TinyLM trunk with a scalar head: V(x, y_<t) for every response position.

    Like Stiennon et al. (2020) and InstructGPT, the critic has its own parameters, so its
    updates can't damage the policy. (They initialize it from the reward model; our RM is
    too simple for that, so we initialize the trunk from the SFT model.)
    """

    def __init__(self, sft_model):
        super().__init__()
        self.trunk = TinyLM()
        self.trunk.load_state_dict(sft_model.state_dict())
        self.v_head = nn.Linear(self.trunk.ln.normalized_shape[0], 1)
        nn.init.zeros_(self.v_head.weight)
        nn.init.zeros_(self.v_head.bias)

    def forward(self, out):
        """(B, max_new): value of the state just BEFORE each response token is generated.

        Aligned with `sequence_logprobs`: column k is V(x, y_<k), the state in which y_k was chosen.
        """
        ids, L = out["ids"], out["prompt_len"]
        h = hidden_states(self.trunk, ids[:, :-1])
        return self.v_head(h).squeeze(-1)[:, L - 1:]


# ---------------------------------------------------------------- small utilities
def masked_mean(x, mask):
    return (x * mask).sum() / mask.sum().clamp(min=1)


def masked_whiten(x, mask, eps=1e-8):
    """Normalize x to mean 0, std 1 over the positions where mask == 1 (others set to 0)."""
    mean = masked_mean(x, mask)
    var = masked_mean((x - mean) ** 2, mask)
    return (x - mean) / torch.sqrt(var + eps) * mask


def sft_loss(model, batch_size, gen):
    """Next-token cross-entropy on a fresh batch of CORRECT additions (the "pretraining" data).

    Used by the PPO-ptx option: mixing this loss into RLHF keeps the skill the reward never asked for.
    """
    import torch.nn.functional as F
    from tinylm import V, encode, encode_prompts

    prompts, answers = make_batch(batch_size, gen)
    ids, L = encode_prompts(prompts)
    resp = torch.full((batch_size, MAX_ANSWER), PAD, dtype=torch.long)
    for i, a in enumerate(answers):
        toks = encode(a) + [EOS]
        resp[i, : len(toks)] = torch.tensor(toks)
    full = torch.cat([ids, resp], 1)
    logits = model(full[:, :-1])[:, L - 1:]
    return F.cross_entropy(logits.reshape(-1, V), resp.reshape(-1), ignore_index=PAD)