import random
from typing import Tuple
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch import Tensor

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
        super().__init__()

        self.input_dim = input_dim
        self.emb_dim = emb_dim
        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        self.dropout = dropout

        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        embedded = self.dropout(self.embedding(src))
        outputs, hidden = self.rnn(embedded)
        # output of bi-directional rnn should be concatenated
        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
        return outputs, hidden
        
class Attention(nn.Module):
    def __init__(self, enc_hid_dim, dec_hid_dim, attn_dim):
        super().__init__()

        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        self.attn_in = (enc_hid_dim * 2) + dec_hid_dim
        self.attn = nn.Linear(self.attn_in, attn_dim)

    def forward(self, decoder_hidden, encoder_outputs):

        src_len = encoder_outputs.shape[0]
        repeated_decoder_hidden = decoder_hidden.unsqueeze(1).repeat(1, src_len, 1)
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        energy = torch.tanh(self.attn(torch.cat((
            repeated_decoder_hidden,
            encoder_outputs),
            dim = 2)))
        attention = torch.sum(energy, dim=2)
        return F.softmax(attention, dim=1)
        
class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
        super().__init__()

        self.emb_dim = emb_dim
        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        self.output_dim = output_dim
        self.dropout = dropout
        self.attention = attention

        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
        self.out = nn.Linear(self.attention.attn_in + emb_dim, output_dim)
        self.dropout = nn.Dropout(dropout)


    def _weighted_encoder_rep(self, decoder_hidden, encoder_outputs):

        a = self.attention(decoder_hidden, encoder_outputs)
        a = a.unsqueeze(1)
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        weighted_encoder_rep = torch.bmm(a, encoder_outputs)
        weighted_encoder_rep = weighted_encoder_rep.permute(1, 0, 2)
        return weighted_encoder_rep


    def forward(self, input, decoder_hidden, encoder_outputs):

        input = input.unsqueeze(0)
        embedded = self.dropout(self.embedding(input))
        weighted_encoder_rep = self._weighted_encoder_rep(decoder_hidden, 
                                                          encoder_outputs)
        rnn_input = torch.cat((embedded, weighted_encoder_rep), dim = 2)
        output, decoder_hidden = self.rnn(rnn_input, decoder_hidden.unsqueeze(0))
        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        weighted_encoder_rep = weighted_encoder_rep.squeeze(0)
        output = self.out(torch.cat((output,
                                     weighted_encoder_rep,
                                     embedded), dim = 1))

        return output, decoder_hidden.squeeze(0)
        
        
class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, trg, teacher_forcing_ratio=0.5):

        batch_size = src.shape[1]
        max_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)
        encoder_outputs, hidden = self.encoder(src)
        # first input to the decoder is the <sos> token
        output = trg[0,:]

        for t in range(1, max_len):
            output, hidden = self.decoder(output, hidden, encoder_outputs)
            outputs[t] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.max(1)[1]
            output = (trg[t] if teacher_force else top1)

        return outputs

[((19, 25, 19, 57), 'torch.nn.Embedding', 'nn.Embedding', ({(19, 38, 19, 47): 'input_dim', (19, 49, 19, 56): 'emb_dim'}, {}), '(input_dim, emb_dim)', True, 'import torch.nn as nn\n'), ((20, 19, 20, 69), 'torch.nn.GRU', 'nn.GRU', (), '', True, 'import torch.nn as nn\n'), ((21, 18, 21, 57), 'torch.nn.Linear', 'nn.Linear', ({(21, 28, 21, 43): 'enc_hid_dim * 2', (21, 45, 21, 56): 'dec_hid_dim'}, {}), '(enc_hid_dim * 2, dec_hid_dim)', True, 'import torch.nn as nn\n'), ((22, 23, 22, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(22, 34, 22, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((38, 20, 38, 53), 'torch.nn.Linear', 'nn.Linear', ({(38, 30, 38, 42): 'self.attn_in', (38, 44, 38, 52): 'attn_dim'}, {}), '(self.attn_in, attn_dim)', True, 'import torch.nn as nn\n'), ((49, 20, 49, 44), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((50, 15, 50, 42), 'torch.nn.functional.softmax', 'F.softmax', (), '', True, 'import torch.nn.functional as F\n'), ((63, 25, 63, 58), 'torch.nn.Embedding', 'nn.Embedding', ({(63, 38, 63, 48): 'output_dim', (63, 50, 63, 57): 'emb_dim'}, {}), '(output_dim, emb_dim)', True, 'import torch.nn as nn\n'), ((64, 19, 64, 67), 'torch.nn.GRU', 'nn.GRU', ({(64, 26, 64, 53): 'enc_hid_dim * 2 + emb_dim', (64, 55, 64, 66): 'dec_hid_dim'}, {}), '(enc_hid_dim * 2 + emb_dim, dec_hid_dim)', True, 'import torch.nn as nn\n'), ((65, 19, 65, 74), 'torch.nn.Linear', 'nn.Linear', ({(65, 29, 65, 61): 'self.attention.attn_in + emb_dim', (65, 63, 65, 73): 'output_dim'}, {}), '(self.attention.attn_in + emb_dim, output_dim)', True, 'import torch.nn as nn\n'), ((66, 23, 66, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(66, 34, 66, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((74, 31, 74, 60), 'torch.bmm', 'torch.bmm', ({(74, 41, 74, 42): 'a', (74, 44, 74, 59): 'encoder_outputs'}, {}), '(a, encoder_outputs)', False, 'import torch\n'), ((85, 20, 85, 72), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((90, 26, 92, 56), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((28, 36, 28, 88), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((45, 38, 48, 20), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((109, 18, 109, 66), 'torch.zeros', 'torch.zeros', ({(109, 30, 109, 37): 'max_len', (109, 39, 109, 49): 'batch_size', (109, 51, 109, 65): 'trg_vocab_size'}, {}), '(max_len, batch_size, trg_vocab_size)', False, 'import torch\n'), ((117, 28, 117, 43), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n')]