import torch
import torch.nn as nn

class Attention(nn.Module):
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.W1 = nn.Linear(hidden_size, hidden_size)
        self.W2 = nn.Linear(hidden_size, hidden_size)
        self.v = nn.Linear(hidden_size, 1, bias=False)

    def forward(self, hidden, encoder_outputs):
        sequence_len = encoder_outputs.shape[1]
        hidden = hidden.unsqueeze(1).repeat(1, sequence_len, 1)
        
        energy = torch.tanh(self.W1(encoder_outputs) + self.W2(hidden))
        attention = self.v(energy).squeeze(2)
        attention_weights = torch.softmax(attention, dim=1)
        
        context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        return context, attention_weights

class SimpleRecurrentNetworkWithAttention(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, cell_type='RNN', device='cpu'):
        super(SimpleRecurrentNetworkWithAttention, self).__init__()
        
        self.device = device
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.attention = Attention(hidden_size * 2)
        
        if cell_type == 'LSTM':
            self.rnn = nn.LSTM(hidden_size, hidden_size, batch_first=True, bidirectional=True)
        elif cell_type == 'GRU':
            self.rnn = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
        else:
            self.rnn = nn.RNN(hidden_size, hidden_size, batch_first=True, bidirectional=True)
            
        self.fc = nn.Linear(hidden_size * 2, output_size)

    def forward(self, inputs):
        embedded = self.embedding(inputs.to(self.device))
        rnn_output, hidden = self.rnn(embedded)
        
        if isinstance(hidden, tuple):
            hidden = hidden[0]
        
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        context, attention_weights = self.attention(hidden, rnn_output)
        output = self.fc(context)
        
        return output, attention_weights