class TinyTransformerClassifier(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int = 64,
        num_heads: int = 4,
        d_ff: int = 128,
        num_layers: int = 2,
        max_len: int = 20,
    ):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos_enc = PositionalEncoding(d_model, max_len)
        self.layers = nn.ModuleList([
            TransformerEncoderBlock(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ])
        self.classifier = nn.Linear(d_model, 2)  # sorted / not sorted

    def forward(self, x):
        h = self.pos_enc(self.embed(x))
        for layer in self.layers:
            h = layer(h)
        pooled = h.mean(dim=1)  # mean-pool over sequence
        return self.classifier(pooled)


def make_batch(batch_size, seq_len, vocab_size, device="cpu"):
    seqs, labels = [], []
    for _ in range(batch_size):
        vals = torch.randint(0, vocab_size, (seq_len,))
        if torch.rand(1).item() < 0.5:
            vals, _ = torch.sort(vals)
            labels.append(1)
        else:
            labels.append(0)
        seqs.append(vals)
    return torch.stack(seqs).to(device), torch.tensor(labels).to(device)


if __name__ == "__main__":
    torch.manual_seed(0)
    vocab_size, seq_len = 20, 10
    model = TinyTransformerClassifier(vocab_size, max_len=seq_len)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for step in range(500):
        x, y = make_batch(64, seq_len, vocab_size)
        logits = model(x)
        loss = F.cross_entropy(logits, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 100 == 0:
            acc = (logits.argmax(-1) == y).float().mean().item()
            msg = f"step {step:4d} loss {loss.item():.4f} acc {acc:.3f}"
            print(msg)
