## Let's build a simple transformer with character level inputs and next character prediction
import torch
import torch.nn as nn
import torch.nn.functional as F


class Transformer(nn.Module):
    class Block(nn.Module):
        def __init__(self, n_embd, n_head, block_size, dropout):
            super().__init__()
            assert n_embd % n_head == 0
            self.attn = nn.MultiheadAttention(
                n_embd, n_head, dropout=dropout, batch_first=True
            )
            self.ln1 = nn.LayerNorm(n_embd)
            self.ff = nn.Sequential(
                nn.Linear(n_embd, 4 * n_embd),
                nn.ReLU(),
                nn.Linear(4 * n_embd, n_embd),
                nn.Dropout(dropout),
            )
            self.ln2 = nn.LayerNorm(n_embd)

        def forward(self, x):
            _, T, _ = x.size()
            # Use generate_square_subsequent_mask to create the mask
            attn_mask = nn.Transformer.generate_square_subsequent_mask(T).to(x.device)
            attn_output, _ = self.attn(x, x, x, is_causal=True, attn_mask=attn_mask)
            x = x + attn_output
            x = self.ln1(x)
            x = x + self.ff(x)
            x = self.ln2(x)
            return x

    def __init__(
        self, vocab_size, n_embd=384, n_head=6, n_layer=6, block_size=128, dropout=0
    ):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(
            *[self.Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]
        )
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)
        self.block_size = block_size
        # self.dropout = nn.Dropout(dropout)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.block_size, "Cannot forward, model block size is exhausted."

        token_embeddings = self.token_embedding(idx)
        position_embeddings = self.position_embedding(
            torch.arange(T, device=idx.device).unsqueeze(0)
        )
        x = token_embeddings + position_embeddings

        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)

        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            return logits, loss
        return logits


# Transformer(vocab_size=65)  # Example instantiation with vocab size of 65
