import argparse
import math

import torch
from torch import nn

from train_bigram import START, STOP, TEST, TRAIN, UNKNOWN, tokenize

CONTEXT = 16
WIDTH = 48


class TinyDecoder(nn.Module):
    def __init__(self, vocabulary_size):
        super().__init__()
        self.tokens = nn.Embedding(vocabulary_size, WIDTH)
        self.positions = nn.Embedding(CONTEXT, WIDTH)
        self.block = nn.TransformerEncoderLayer(
            d_model=WIDTH, nhead=4, dim_feedforward=96, dropout=0, batch_first=True
        )
        self.output = nn.Linear(WIDTH, vocabulary_size)

    def forward(self, ids):
        positions = torch.arange(ids.shape[1], device=ids.device)
        hidden = self.tokens(ids) + self.positions(positions)
        future = positions[None, :] > positions[:, None]
        return self.output(self.block(hidden, src_mask=future))


def examples(documents, ids):
    rows = []
    for line in documents:
        tokens = [START, *tokenize(line), STOP]
        numbers = [ids.get(token, ids[UNKNOWN]) for token in tokens]
        if len(numbers) - 1 > CONTEXT:
            raise ValueError(f"Keep each training or test line within {CONTEXT} tokens plus its end marker.")
        rows.append(numbers)
    length = max(len(row) - 1 for row in rows)
    # Use the start ID as input padding. Ignore padded targets in the loss.
    inputs = [row[:-1] + [ids[START]] * (length + 1 - len(row)) for row in rows]
    targets = [row[1:] + [-100] * (length + 1 - len(row)) for row in rows]
    return torch.tensor(inputs), torch.tensor(targets)


def next_token_loss(model, inputs, targets):
    logits = model(inputs)
    return nn.functional.cross_entropy(logits.flatten(0, 1), targets.flatten())


@torch.inference_mode()
def show_scores(model, train, test, step):
    train_score = math.exp(next_token_loss(model, *train).item())
    test_score = math.exp(next_token_loss(model, *test).item())
    print(f"Step {step:3} | train perplexity {train_score:.2f} | test perplexity {test_score:.2f}")


@torch.inference_mode()
def generate(model, vocabulary, ids, prefix):
    history = [START, *tokenize(prefix)]
    while len(history) <= CONTEXT:
        inputs = torch.tensor([[ids.get(word, ids[UNKNOWN]) for word in history]])
        logits = model(inputs)[0, -1]
        logits[ids[START]] = -torch.inf
        token = vocabulary[logits.argmax().item()]
        if token == STOP:
            break
        history.append(token)
    return " ".join(history[1:])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a tiny causal Transformer on the CPU.")
    parser.add_argument("--steps", type=int, default=120)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--prefix", default="the library opens at")
    args = parser.parse_args()
    if args.steps < 1 or len(tokenize(args.prefix)) >= CONTEXT:
        parser.error(f"Use a positive step count and a prefix shorter than {CONTEXT} tokens.")
    torch.manual_seed(args.seed)
    torch.set_num_threads(1)
    vocabulary = sorted({START, STOP, UNKNOWN, *(token for line in TRAIN for token in tokenize(line))})
    ids = {token: index for index, token in enumerate(vocabulary)}
    train_data = examples(TRAIN, ids)
    test_data = examples(TEST, ids)
    model = TinyDecoder(len(vocabulary))
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003)
    print(f"Trainable parameters: {sum(item.numel() for item in model.parameters()):,}")
    show_scores(model, train_data, test_data, 0)
    for step in range(1, args.steps + 1):
        optimizer.zero_grad()
        loss = next_token_loss(model, *train_data)
        loss.backward()
        optimizer.step()
        if step % 30 == 0 or step == args.steps:
            show_scores(model, train_data, test_data, step)
    print("Greedy completion:", generate(model, vocabulary, ids, args.prefix))