import torch

class SimpleTokenizer:
    def __init__(self, vocab):
        self.stoi = vocab
        self.unk = vocab.get("<unk>", 1)

    def __call__(self, text, return_tensors="pt", max_length=256, **kwargs):
        tokens = text.lower().split()
        ids = [self.stoi.get(t, self.unk) for t in tokens][:max_length]
        return {"input_ids": torch.tensor([ids], dtype=torch.long)}

class LSTMClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes):
        super().__init__()

        self.embedding = torch.nn.Embedding(vocab_size, embedding_dim)

        self.lstm = torch.nn.LSTM(
            embedding_dim,
            hidden_dim,
            bidirectional=True,
            batch_first=True
        )

        self.fc = torch.nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        _, (hidden, _) = self.lstm(x)
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        return self.fc(hidden)


class GRUClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes):
        super().__init__()

        self.embedding = torch.nn.Embedding(vocab_size, embedding_dim)

        self.gru = torch.nn.GRU(
            embedding_dim,
            hidden_dim,
            batch_first=True
        )

        self.fc = torch.nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        _, hidden = self.gru(x)
        return self.fc(hidden[-1])


class BiLSTMClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes):
        super().__init__()

        self.embedding = torch.nn.Embedding(vocab_size, embedding_dim)

        self.lstm = torch.nn.LSTM(
            embedding_dim,
            hidden_dim,
            bidirectional=True,
            batch_first=True
        )

        self.fc = torch.nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        _, (hidden, _) = self.lstm(x)
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        return self.fc(hidden)


class BiGRUClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes):
        super().__init__()

        self.embedding = torch.nn.Embedding(vocab_size, embedding_dim)

        self.gru = torch.nn.GRU(
            embedding_dim,
            hidden_dim,
            bidirectional=True,
            batch_first=True
        )

        self.fc = torch.nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        _, hidden = self.gru(x)
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        return self.fc(hidden)