from collections import Counter
import torch
import torch.nn as nn


class Vocabulary:
    def __init__(self, min_freq=1):
        self.min_freq = min_freq
        self.token2id = {'<PAD>': 0, '<UNK>': 1}
        self.id2token = {0: '<PAD>', 1: '<UNK>'}

    def build(self, texts):
        counter = Counter(tok for text in texts for tok in text.split())
        for token, freq in counter.items():
            if freq >= self.min_freq and token not in self.token2id:
                idx = len(self.token2id)
                self.token2id[token] = idx
                self.id2token[idx] = token
        return self

    def encode(self, text, max_len=None):
        ids = [self.token2id.get(t, 1) for t in text.split()]
        if max_len:
            ids = ids[:max_len]
        return ids if ids else [1]

    def __len__(self):
        return len(self.token2id)


class AttentionPooling(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.attn = nn.Linear(hidden_dim, 1)

    def forward(self, h, mask=None):
        scores = self.attn(h).squeeze(-1)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        weights = torch.softmax(scores, dim=-1)
        return (h * weights.unsqueeze(-1)).sum(dim=1)


class BiLSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, num_classes, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0.0,
        )
        self.attention  = AttentionPooling(hidden_dim * 2)
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_dim * 2, num_classes),
        )

    def forward(self, input_ids, attention_mask=None):
        x         = self.embedding(input_ids)
        output, _ = self.lstm(x)
        pooled    = self.attention(output, attention_mask)
        return self.classifier(pooled)