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

[((494, 526), 'torch.nn.Embedding', 'nn.Embedding', (['input_dim', 'emb_dim'], {}), '(input_dim, emb_dim)\n', (506, 526), True, 'import torch.nn as nn\n'), ((546, 594), 'torch.nn.GRU', 'nn.GRU', (['emb_dim', 'enc_hid_dim'], {'bidirectional': '(True)'}), '(emb_dim, enc_hid_dim, bidirectional=True)\n', (552, 594), True, 'import torch.nn as nn\n'), ((615, 654), 'torch.nn.Linear', 'nn.Linear', (['(enc_hid_dim * 2)', 'dec_hid_dim'], {}), '(enc_hid_dim * 2, dec_hid_dim)\n', (624, 654), True, 'import torch.nn as nn\n'), ((678, 697), 'torch.nn.Dropout', 'nn.Dropout', (['dropout'], {}), '(dropout)\n', (688, 697), True, 'import torch.nn as nn\n'), ((1287, 1320), 'torch.nn.Linear', 'nn.Linear', (['self.attn_in', 'attn_dim'], {}), '(self.attn_in, attn_dim)\n', (1296, 1320), True, 'import torch.nn as nn\n'), ((1725, 1749), 'torch.sum', 'torch.sum', (['energy'], {'dim': '(2)'}), '(energy, dim=2)\n', (1734, 1749), False, 'import torch\n'), ((1765, 1792), 'torch.nn.functional.softmax', 'F.softmax', (['attention'], {'dim': '(1)'}), '(attention, dim=1)\n', (1774, 1792), True, 'import torch.nn.functional as F\n'), ((2185, 2218), 'torch.nn.Embedding', 'nn.Embedding', (['output_dim', 'emb_dim'], {}), '(output_dim, emb_dim)\n', (2197, 2218), True, 'import torch.nn as nn\n'), ((2238, 2284), 'torch.nn.GRU', 'nn.GRU', (['(enc_hid_dim * 2 + emb_dim)', 'dec_hid_dim'], {}), '(enc_hid_dim * 2 + emb_dim, dec_hid_dim)\n', (2244, 2284), True, 'import torch.nn as nn\n'), ((2306, 2361), 'torch.nn.Linear', 'nn.Linear', (['(self.attention.attn_in + emb_dim)', 'output_dim'], {}), '(self.attention.attn_in + emb_dim, output_dim)\n', (2315, 2361), True, 'import torch.nn as nn\n'), ((2385, 2404), 'torch.nn.Dropout', 'nn.Dropout', (['dropout'], {}), '(dropout)\n', (2395, 2404), True, 'import torch.nn as nn\n'), ((2655, 2684), 'torch.bmm', 'torch.bmm', (['a', 'encoder_outputs'], {}), '(a, encoder_outputs)\n', (2664, 2684), False, 'import torch\n'), ((3116, 3166), 'torch.cat', 'torch.cat', (['(embedded, weighted_encoder_rep)'], {'dim': '(2)'}), '((embedded, weighted_encoder_rep), dim=2)\n', (3125, 3166), False, 'import torch\n'), ((3414, 3472), 'torch.cat', 'torch.cat', (['(output, weighted_encoder_rep, embedded)'], {'dim': '(1)'}), '((output, weighted_encoder_rep, embedded), dim=1)\n', (3423, 3472), False, 'import torch\n'), ((923, 981), 'torch.cat', 'torch.cat', (['(hidden[(-2), :, :], hidden[(-1), :, :])'], {'dim': '(1)'}), '((hidden[(-2), :, :], hidden[(-1), :, :]), dim=1)\n', (932, 981), False, 'import torch\n'), ((1603, 1663), 'torch.cat', 'torch.cat', (['(repeated_decoder_hidden, encoder_outputs)'], {'dim': '(2)'}), '((repeated_decoder_hidden, encoder_outputs), dim=2)\n', (1612, 1663), False, 'import torch\n'), ((4006, 4054), 'torch.zeros', 'torch.zeros', (['max_len', 'batch_size', 'trg_vocab_size'], {}), '(max_len, batch_size, trg_vocab_size)\n', (4017, 4054), False, 'import torch\n'), ((4377, 4392), 'random.random', 'random.random', ([], {}), '()\n', (4390, 4392), False, 'import random\n')]