r"""Helper code for Lab 18 (complete: you do not need to edit this file).

What is in here:
  * the token vocabulary of the toy agent (digits, tool-call words, tool-output words),
  * Policy:  a tiny GRU language model over that vocabulary (PyTorch, the "trainer"),
  * Sampler: a numpy copy of the same network for fast token-by-token sampling
             (the "inference engine"; Chapter 20 explains why real systems split these),
  * Trajectory: the record of one multi-turn episode,
  * warm_start: a short supervised "cold start" that teaches the tool-call format,
  * make_batch: pads trajectories into tensors and builds the loss mask,
  * render: pretty-prints a trajectory with each token's origin.

A trajectory is one flat token sequence, exactly like an LLM agent's context window:

    <bos> ?  call 3  <tool> higher  call 5  <tool> lower  ans 4
    \_prompt/ \model/ \___tool____/ \model/ \___tool___/ \model/

Prompt tokens and tool tokens are written by the environment; only the model tokens
were sampled from the policy.
"""

from dataclasses import dataclass, field

import numpy as np
import torch
import torch.nn as nn

torch.set_num_threads(1)

# ---------------------------------------------------------------- vocabulary
N_NUMBERS = 8                                   # the secret is one of 0..7
NAMES = [str(d) for d in range(N_NUMBERS)] + [
    "call",     # model: start a tool call  compare(x)  ->  "call x"
    "ans",      # model: give the final answer          ->  "ans x"
    "lower",    # tool output: the secret is lower than x
    "higher",   # tool output: the secret is higher than x
    "equal",    # tool output: the secret equals x
    "<tool>",   # tool output: marks the start of a tool result
    "<bos>",    # prompt: beginning of an episode
    "?",        # prompt: "the secret is hidden" (a hard task)
    "<pad>",    # padding, never part of an episode
]
TOK = {name: i for i, name in enumerate(NAMES)}
CALL, ANS = TOK["call"], TOK["ans"]
LOWER, HIGHER, EQUAL = TOK["lower"], TOK["higher"], TOK["equal"]
TOOL, BOS, HIDDEN, PAD = TOK["<tool>"], TOK["<bos>"], TOK["?"], TOK["<pad>"]
V = len(NAMES)
TURN_LEN = 2                                    # every model turn is two tokens: verb, number


def is_number(tok):
    return 0 <= tok < N_NUMBERS


# ---------------------------------------------------------------- one episode
@dataclass
class Trajectory:
    """One multi-turn episode, as a flat token sequence plus bookkeeping.

    tokens:       every token of the episode in order (prompt, model turns, tool outputs).
    is_model:     1 where the token was sampled from the policy, 0 where the environment wrote it
                  (prompt tokens and tool outputs). Same length as `tokens`.
    turn_id:      which turn each token belongs to: -1 for the prompt, k for model turn k and
                  the tool output that answered it. Same length as `tokens`.
    turn_rewards: one reward per model turn (length = number of turns).
    """
    tokens: list = field(default_factory=list)
    is_model: list = field(default_factory=list)
    turn_id: list = field(default_factory=list)
    turn_rewards: list = field(default_factory=list)

    @property
    def total_reward(self):
        return float(sum(self.turn_rewards))

    @property
    def n_turns(self):
        return len(self.turn_rewards)

    @property
    def n_calls(self):
        return sum(1 for i, t in enumerate(self.tokens) if t == CALL and self.is_model[i])

    @property
    def solved(self):
        """True if the episode ended with a correct answer (final turn reward is +1)."""
        return bool(self.turn_rewards) and self.turn_rewards[-1] >= 1.0


def render(traj):
    """A readable one-line view: [prompt] model tokens {tool output} ..."""
    wrap = {"prompt": "[{}]", "model": "{}", "tool": "{{{}}}"}
    parts, kind, buf = [], None, []
    for tok, m, k in zip(traj.tokens, traj.is_model, traj.turn_id):
        new = "prompt" if k < 0 else ("model" if m else "tool")
        if new != kind and buf:
            parts.append(wrap[kind].format(" ".join(buf)))
            buf = []
        kind = new
        buf.append(NAMES[tok])
    if buf:
        parts.append(wrap[kind].format(" ".join(buf)))
    return "  ".join(parts) + f"    turn rewards {traj.turn_rewards}"


def sample_token(probs, rng):
    """Draw one token id from a probability vector (faster than rng.choice)."""
    return int(min(np.searchsorted(np.cumsum(probs), rng.random() * probs.sum()), len(probs) - 1))


# ---------------------------------------------------------------- the policy
class Policy(nn.Module):
    """A tiny recurrent language model: embedding -> GRU -> linear head over V tokens.

    forward(ids) takes (B, T) token ids and returns (B, T, V) logits, where position t
    predicts token t+1 (the same convention as tinylm.py in Labs 11-17).

    The logits of <bos> and <pad> are set to -1e9 (probability 0), here and in the Sampler, so the trainer
    and the sampler agree on which tokens can be generated. (If only the sampler forbade
    them, the trainer's softmax would quietly move probability onto tokens that are never
    sampled: a small train-inference mismatch of the kind Chapter 20 discusses.)
    """

    def __init__(self, d_emb=32, d_hidden=64):
        super().__init__()
        self.emb = nn.Embedding(V, d_emb)
        self.gru = nn.GRU(d_emb, d_hidden, batch_first=True)
        self.head = nn.Linear(d_hidden, V)

    def forward(self, ids):
        h, _ = self.gru(self.emb(ids))
        logits = self.head(h)
        never = torch.zeros(V, dtype=torch.bool)
        never[[BOS, PAD]] = True
        return logits.masked_fill(never, -1e9)

    def sampler(self):
        """A frozen numpy snapshot of the current weights, for fast rollouts."""
        return Sampler(self)


class Sampler:
    """The same network as Policy, in numpy, one token at a time.

    Usage:
        state = sampler.start()
        probs, state = sampler.step(state, token)   # feed a token, get next-token probabilities

    Every token in the episode (prompt, model and tool) must be fed with step(), because
    the policy has to *read* the tool's output even though it did not write it.
    """

    def __init__(self, policy):
        with torch.no_grad():
            E = policy.emb.weight.numpy().astype(np.float64)
            W_ih = policy.gru.weight_ih_l0.numpy().astype(np.float64)
            b_ih = policy.gru.bias_ih_l0.numpy().astype(np.float64)
            self.W_hh = policy.gru.weight_hh_l0.numpy().astype(np.float64)
            self.b_hh = policy.gru.bias_hh_l0.numpy().astype(np.float64)
            self.W_o = policy.head.weight.numpy().astype(np.float64)
            self.b_o = policy.head.bias.numpy().astype(np.float64)
        self.gi = E @ W_ih.T + b_ih          # (V, 3H): input projection of every token
        self.H = self.W_hh.shape[1]

    def start(self):
        return np.zeros(self.H)

    def step(self, h, tok):
        H = self.H
        gi, gh = self.gi[tok], self.W_hh @ h + self.b_hh
        r = 1.0 / (1.0 + np.exp(-(gi[:H] + gh[:H])))
        z = 1.0 / (1.0 + np.exp(-(gi[H:2 * H] + gh[H:2 * H])))
        n = np.tanh(gi[2 * H:] + r * gh[2 * H:])
        h = (1.0 - z) * n + z * h
        logits = self.W_o @ h + self.b_o
        logits[[BOS, PAD]] = -1e9            # never generate these (same as Policy.forward)
        p = np.exp(logits - logits.max())
        return p / p.sum(), h


# ---------------------------------------------------------------- cold start
def _cold_start_episode(rng):
    """A random but well-formed episode: a few random tool calls (stopping if the tool
    says 'equal'), then an answer (the number the tool confirmed, else a random guess).
    It teaches the format and nothing about strategy: the calls are random, and the
    revealed number in the prompt is ignored."""
    secret = int(rng.integers(N_NUMBERS))
    tokens = [BOS, secret if rng.random() < 0.5 else HIDDEN]
    is_model = [0, 0]
    found = None
    for x in rng.permutation(N_NUMBERS)[: int(rng.integers(5))]:
        x = int(x)
        res = EQUAL if x == secret else (LOWER if secret < x else HIGHER)
        tokens += [CALL, x, TOOL, res]
        is_model += [1, 1, 0, 0]
        if res == EQUAL:
            found = x
            break
    tokens += [ANS, found if found is not None else int(rng.integers(N_NUMBERS))]
    is_model += [1, 1]
    return tokens, is_model


def warm_start(policy, steps=500, batch_size=64, lr=3e-3, seed=0):
    """Supervised 'cold start' on random well-formed episodes (like the cold-start SFT
    stage before agentic RL). Trains only on the model's tokens (environment tokens are
    masked here too). Takes a few seconds."""
    rng = np.random.default_rng(seed + 1000)
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    for _ in range(steps):
        eps = [_cold_start_episode(rng) for _ in range(batch_size)]
        L = max(len(t) for t, _ in eps)
        ids = torch.full((batch_size, L), PAD, dtype=torch.long)
        m = torch.zeros(batch_size, L)
        for i, (t, mm) in enumerate(eps):
            ids[i, : len(t)] = torch.tensor(t)
            m[i, : len(mm)] = torch.tensor(mm, dtype=torch.float)
        logp = torch.log_softmax(policy(ids[:, :-1]), -1).gather(-1, ids[:, 1:, None]).squeeze(-1)
        loss = -(logp * m[:, 1:]).sum() / m[:, 1:].sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return policy


# ---------------------------------------------------------------- batching
def make_batch(trajs, turn_advs, mask_env=True):
    """Pad a list of trajectories into training tensors.

    Args:
        trajs:     list of B Trajectory objects.
        turn_advs: list of B arrays; turn_advs[i][k] is the advantage of turn k of trajs[i].
        mask_env:  if True (correct), the loss mask is is_model: only sampled tokens are trained.
                   If False (the bug this lab studies), tool-output tokens are trained too
                   (the prompt and padding are still excluded).

    Returns:
        ids        (B, T) long, padded with <pad>
        loss_mask  (B, T) float, 1 for tokens that should receive gradient
        token_adv  (B, T) float, each token gets the advantage of the turn it belongs to
    """
    B, T = len(trajs), max(len(t.tokens) for t in trajs)
    ids = torch.full((B, T), PAD, dtype=torch.long)
    loss_mask = torch.zeros(B, T)
    token_adv = torch.zeros(B, T)
    for i, (tr, adv) in enumerate(zip(trajs, turn_advs)):
        n = len(tr.tokens)
        ids[i, :n] = torch.tensor(tr.tokens)
        turn = np.array(tr.turn_id)
        in_episode = turn >= 0                                  # not the prompt
        keep = np.array(tr.is_model, dtype=bool) if mask_env else in_episode
        loss_mask[i, :n] = torch.tensor(keep, dtype=torch.float)
        a = np.zeros(n)
        a[in_episode] = np.asarray(adv, dtype=float)[turn[in_episode]]
        token_adv[i, :n] = torch.tensor(a, dtype=torch.float)
    return ids, loss_mask, token_adv