"""A tiny causal transformer language model and a verifiable arithmetic task.

Small enough to train on a laptop CPU in well under a minute, big enough to show
real LLM post-training dynamics (SFT, then RL with a checkable reward).

The task: given a prompt like "37+48=", produce the digits of the sum and stop.
    prompt tokens:   <bos> 3 7 + 4 8 =
    response tokens: 8 5 <eos>
The reward is verifiable: 1.0 if the response is exactly the right answer, else 0.0.

Typical use:
    torch.manual_seed(0)
    model = TinyLM()
    sft(model, steps=300)                       # imitation: a partly-competent starting policy
    prompts, answers = make_batch(16, rng)      # a batch of problems
    out = sample(model, prompts, temperature=1.0)
    rewards = [reward(r, a) for r, a in zip(decode_responses(out), answers)]
    logp = sequence_logprobs(model, out)        # per-token log-probs of the response, with a mask
"""

import math
import random

import torch
import torch.nn as nn
import torch.nn.functional as F

# Tiny models run several times faster on one CPU thread than on many.
torch.set_num_threads(1)

# ---------------------------------------------------------------- vocabulary
VOCAB = list("0123456789+=") + ["<bos>", "<eos>", "<pad>"]
STOI = {s: i for i, s in enumerate(VOCAB)}
BOS, EOS, PAD = STOI["<bos>"], STOI["<eos>"], STOI["<pad>"]
V = len(VOCAB)
MAX_ANSWER = 4  # answer digits + <eos> fit in this many tokens for sums < 1000


def encode(s):
    return [STOI[c] for c in s]


def decode(ids):
    out = []
    for i in ids:
        if i == EOS:
            break
        if i in (BOS, PAD):
            continue
        out.append(VOCAB[i])
    return "".join(out)


# ---------------------------------------------------------------- task
def make_problem(rng, max_n=99):
    a, b = rng.randint(0, max_n), rng.randint(0, max_n)
    return f"{a}+{b}=", str(a + b)


def make_batch(n, rng, max_n=99):
    """Return (prompts, answers): lists of strings like ('37+48=', '85')."""
    probs = [make_problem(rng, max_n) for _ in range(n)]
    return [p for p, _ in probs], [a for _, a in probs]


def reward(response, answer):
    """Verifiable reward: 1.0 for an exactly correct answer, else 0.0."""
    return 1.0 if response == answer else 0.0


# ---------------------------------------------------------------- model
class Block(nn.Module):
    def __init__(self, d, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))


class TinyLM(nn.Module):
    def __init__(self, d=64, n_head=4, n_layer=2, max_len=16):
        super().__init__()
        self.max_len = max_len
        self.tok = nn.Embedding(V, d)
        self.pos = nn.Embedding(max_len, d)
        self.blocks = nn.ModuleList([Block(d, n_head) for _ in range(n_layer)])
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, V)

    def forward(self, ids):
        """ids: (B, T) long -> logits (B, T, V). Position t predicts token t+1."""
        T = ids.shape[1]
        mask = torch.triu(torch.full((T, T), float("-inf")), diagonal=1)
        x = self.tok(ids) + self.pos(torch.arange(T))
        for blk in self.blocks:
            x = blk(x, mask)
        return self.head(self.ln(x))


# ---------------------------------------------------------------- batching helpers
def encode_prompts(prompts):
    """Left-pad prompts so they all end at the same column. Returns (ids, prompt_len)."""
    seqs = [[BOS] + encode(p) for p in prompts]
    L = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), L), PAD, dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, L - len(s):] = torch.tensor(s)
    return ids, L


@torch.no_grad()
def sample(model, prompts, temperature=1.0, max_new=MAX_ANSWER, greedy=False):
    """Sample responses. Returns a dict:
        ids          (B, L+max_new) prompt + response tokens (response padded with <pad> after <eos>)
        prompt_len   int, column where responses start
        mask         (B, max_new) float, 1 for response tokens that were generated (incl. <eos>)
    """
    model.eval()
    ids, L = encode_prompts(prompts)
    B = ids.shape[0]
    done = torch.zeros(B, dtype=torch.bool)
    mask = torch.zeros(B, max_new)
    for t in range(max_new):
        logits = model(ids)[:, -1, :] / max(temperature, 1e-6)
        logits[:, [BOS, PAD]] = float("-inf")
        nxt = logits.argmax(-1) if greedy else torch.multinomial(F.softmax(logits, -1), 1).squeeze(1)
        nxt = torch.where(done, torch.full_like(nxt, PAD), nxt)
        mask[:, t] = (~done).float()
        ids = torch.cat([ids, nxt[:, None]], dim=1)
        done |= nxt == EOS
    model.train()
    return {"ids": ids, "prompt_len": L, "mask": mask}


def decode_responses(out):
    L = out["prompt_len"]
    return [decode(row[L:].tolist()) for row in out["ids"]]


def sequence_logprobs(model, out):
    """Per-token log-probs of the sampled response tokens under `model` (with grad).

    Returns (B, max_new) tensor; multiply by out["mask"] before summing or averaging.
    """
    ids, L = out["ids"], out["prompt_len"]
    logits = model(ids[:, :-1])  # position j predicts token j+1
    logp = F.log_softmax(logits, -1)
    tgt = ids[:, 1:]
    tok_logp = logp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1)  # (B, T-1)
    return tok_logp[:, L - 1:]  # response tokens only


# ---------------------------------------------------------------- SFT and evaluation
def sft(model, steps=300, batch_size=64, lr=3e-3, seed=0, max_n=99, verbose=False):
    """Supervised fine-tuning (next-token cross-entropy on correct answers).

    A few hundred steps gives a partly competent model: a good starting point for RL.
    """
    rng = random.Random(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    for step in range(steps):
        prompts, answers = make_batch(batch_size, rng, max_n)
        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:]
        loss = F.cross_entropy(logits.reshape(-1, V), resp.reshape(-1), ignore_index=PAD)
        opt.zero_grad()
        loss.backward()
        opt.step()
        if verbose and step % 100 == 0:
            print(f"sft step {step:4d}  loss {loss.item():.3f}")
    return model


def load_or_train_sft(path, steps=350, seed=0):
    """Load a TinyLM checkpoint from `path`, or SFT one from scratch (seeded) and save it.

    ~350 steps gives roughly 30-50% sampled accuracy: room for RL to improve.
    """
    import os

    torch.manual_seed(seed)
    model = TinyLM()
    if os.path.exists(path):
        model.load_state_dict(torch.load(path))
    else:
        sft(model, steps=steps, seed=seed)
        torch.save(model.state_dict(), path)
    return model


@torch.no_grad()
def accuracy(model, n=500, seed=123, max_n=99, greedy=True):
    rng = random.Random(seed)
    prompts, answers = make_batch(n, rng, max_n)
    resp = decode_responses(sample(model, prompts, greedy=greedy))
    return sum(r == a for r, a in zip(resp, answers)) / n


if __name__ == "__main__":
    import time

    torch.manual_seed(0)
    m = TinyLM()
    t0 = time.time()
    for k in range(1, 7):
        sft(m, steps=100, seed=k)
        print(f"after {100*k} SFT steps: greedy acc {accuracy(m):.2f}, "
              f"sampled acc {accuracy(m, greedy=False):.2f}  ({time.time() - t0:.1f}s)")