import torch
import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        
        # Initialize the encoder with required dimensions and parameters
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        
        # Embedding layer to convert input tokens to dense vectors
        self.embedding = nn.Embedding(input_dim, emb_dim)
        
        # Recurrent layer (GRU) to process the embedded input
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout = dropout)
        
        # Dropout layer to prevent overfitting
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        
        # Embedding the source sequence
        embedded = self.dropout(self.embedding(src))
        
        # Forward pass through the GRU
        outputs, hidden = self.rnn(embedded)
        
        return hidden

class BahdanauAttention(nn.Module):
    def __init__(self, hidden_dim):
        super(BahdanauAttention, self).__init__()

        # Initialize the attention mechanism with the hidden dimension
        self.hidden_dim = hidden_dim

        self.attention = nn.Linear(hidden_dim * 2, hidden_dim)  # Linear layer
        self.v = nn.Parameter(torch.rand(hidden_dim))  # Attention scoring

    def forward(self, hidden, encoder_outputs):
        # Get the number of time steps in the encoder outputs
        timestep = encoder_outputs.shape[0]
        h = hidden.repeat(timestep, 1, 1).transpose(0, 1)

        # Prepare encoder outputs for attention calculation
        encoder_outputs = encoder_outputs.transpose(0, 1)

        # Get attention energies
        attn_energies = self.score(h, encoder_outputs)

        # Apply softmax to get attention weights
        return F.softmax(attn_energies, dim=1).unsqueeze(1)

    def score(self, hidden, encoder_outputs):
        # Get hidden states and encoder outputs
        energy = F.tanh(self.attention(torch.cat([hidden, encoder_outputs], 2)))
        energy = energy.transpose(1, 2)

        # Get attention scores
        v = self.v.repeat(encoder_outputs.data.shape[0], 1).unsqueeze(1)
        energy = torch.bmm(v, energy)

        return energy.squeeze(1)

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout, attention):
        super().__init__()

        # Initialize parameters and attention mechanism
        self.output_dim = output_dim
        self.attention = attention
        
        # Embedding layer for output tokens
        self.embedding = nn.Embedding(output_dim, emb_dim)
        
        # Recurrent layer for decoding with attention
        self.rnn = nn.GRU((hid_dim * 2) + emb_dim, hid_dim, n_layers, dropout=dropout)
        
        # Connect layers for final output
        self.fc_out = nn.Linear((hid_dim * 2) + hid_dim + emb_dim, output_dim)
        
        # Dropout layer to prevent overfitting
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, encoder_outputs):
        
        # Add a singleton dimension to input
        input = input.unsqueeze(0)
        
        # Embedding the input
        embedded = self.dropout(self.embedding(input))
        
        # Get attention and context vector
        a = self.attention(hidden, encoder_outputs)                
        a = a.bmm(encoder_outputs.transpose(0, 1))        
        a = a.transpose(0, 1)
        
        # Concatenate embedded input and context vector
        rnn_input = torch.cat((embedded, a), dim = 2)
        
        # Forward pass through the GRU
        output, hidden = self.rnn(rnn_input, hidden)
        
        # Squeeze tensors
        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        a = a.squeeze(0)
        
        # Concatenate the output, context vector, and embedded input
        prediction = self.fc_out(torch.cat((output, a, embedded), dim = 1))
        
        return prediction, hidden

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        
        # Initialize encoder, decoder, and device
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
    def forward(self, src, trg, teacher_forcing_ratio = 0.5):
        
        # Get dimensions
        batch_size = src.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        
        # Store decoder outputs
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)        
        encoder_outputs = self.encoder(src)
        
        # Initialize decoder hidden state with encoder final hidden state
        hidden = encoder_outputs
        
        # First input token
        input = trg[0,:]
        
        for t in range(1, trg_len):

            # Forward pass through the decoder
            output, hidden = self.decoder(input, hidden, encoder_outputs)
            
            # Save decoder output
            outputs[t] = output
            
            teacher_force = random.random() < teacher_forcing_ratio
            
            # Get the next input token
            top1 = output.argmax(1)             
            input = trg[t] if teacher_force else top1

        return outputs