# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_seq2seq.ipynb (unless otherwise specified).

__all__ = ['Encoder', 'NewDecoder', 'Seq2Seq']

# Cell
from torch import nn
from torch import optim
import torch
import torch.nn.functional as F

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Cell
class Encoder(nn.Module):
    def __init__(self, input_size, embedding_size, hidden_size, num_layers=2, p=0.1):
        super(Encoder, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.dropout = nn.Dropout(p)
        self.embedding = nn.Embedding(input_size, embedding_size)
        self.rnn = nn.LSTM(embedding_size, hidden_size, num_layers, dropout=p,batch_first=False)

    def forward(self, x, x_len):
        # x shape (seq_length, N)
        embedding = self.dropout(self.embedding(x))
        # embedding shape : (seq_length, N, embedding_size)

        x_packed = pack_padded_sequence(embedding, x_len.cpu(), batch_first=False, enforce_sorted=False)
        output_packed, (hidden,cell) = self.rnn(x_packed)

        # irrelevant because we are interested only in hidden state
        #output_padded, output_lengths = pad_packed_sequence(output_packed, batch_first=True)

        # output is irrelevant, context vector is important
        return hidden,cell


# Cell
class NewDecoder(nn.Module):
    def __init__(self, hidden_size, embedding_size, output_size, n_layers=1, dropout_p=0.1):
        super(NewDecoder, self).__init__()

        # Define parameters
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers =n_layers
        self.dropout_p = dropout_p

        # Define layers
        self.embedding = nn.Embedding(output_size, embedding_size)
        self.dropout=nn.Dropout(dropout_p)
        self.rnn = nn.LSTM(embedding_size, hidden_size, n_layers, dropout=dropout_p, batch_first=False)
        self.out = nn.Linear(hidden_size, output_size)

    def forward(self, word_input, last_hidden, encoder_outputs):
        # Note that we will only be running forward for a single decoder time step, but will
        # use all encoder outputs
        word_input = word_input.unsqueeze(0)
        # we are not using encoder_outputs here
        word_embedded = self.embedding(word_input) # 1 X B
        word_embedded = self.dropout(word_embedded)  # 1 X B X emb_length

        # Combine embedded input word and hidden vector, run through RNN
        output, hidden = self.rnn(word_embedded, last_hidden)  # 1 X B X hidden
        predictions = self.out(output)     # 1, B, out
        #output = F.log_softmax(predictions)

        return predictions, hidden

# Cell
import random

import pytorch_lightning as pl
import pytorch_lightning.metrics.functional as plfunc
from pytorch_lightning.loggers import TensorBoardLogger

# Cell
class Seq2Seq(pl.LightningModule):
    """ Encoder decoder pytorch lightning module for training seq2seq model with teacher forcing
    Module try to learn mapping from one sequence to another
    """

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument("--emb_dim", type=int, default=32)
        parser.add_argument('--hidden_dim', type=int, default=64)
        parser.add_argument('--dropout', type=float, default=0.1)

        return parser

    def __init__(self,
                input_vocab_size,
                output_vocab_size,
                padding_index = 0,
                emb_dim = 8,
                hidden_dim=32,
                dropout=0.1,
                max_length=20,
                **kwargs):
        super().__init__()


        # dynamic, based on tokenizer vocab size defined in datamodule
        self.input_dim = input_vocab_size
        self.output_dim = output_vocab_size

        self.enc_emb_dim = emb_dim
        self.dec_emb_dim = emb_dim

        self.enc_hid_dim = hidden_dim
        self.dec_hid_dim = hidden_dim

        self.enc_dropout = dropout
        self.dec_dropout = dropout

        self.pad_idx = padding_index
        self.num_layers = 2
        self.max_length =10
        self.save_hyperparameters()

        self.max_epochs= kwargs.get('max_epochs',5)

        self.learning_rate = 0.0005

        self._loss = nn.CrossEntropyLoss(ignore_index=self.pad_idx)

        self.encoder = Encoder(
            self.input_dim,
            self.enc_emb_dim,
            self.enc_hid_dim,
            self.num_layers,
            self.enc_dropout
        )

        self.decoder = NewDecoder(
            self.enc_hid_dim,
            self.dec_emb_dim,
            self.output_dim,
            self.num_layers,
            self.dec_dropout
        )

        self._init_weights()

    def _init_weights(self):
        for name, param in self.named_parameters():
            if "weight" in name:
                nn.init.normal_(param.data, mean=0, std=0.01)

            else:
                nn.init.constant_(param.data, 0)

    def create_mask(self, src):
        mask = (src != self.pad_idx).permute(1, 0)
        return mask



    def forward(self, src_seq, source_len, trg_seq, teacher_force_ratio=0.5):
        """
        teacher_force_ratio is used to help in decoding.
        In starting, original input token will be sent as input token
        """
        source = src_seq.transpose(0, 1)
        target_len = self.max_length

        if trg_seq is not None:
            target = trg_seq.transpose(0, 1)
            target_len = target.shape[0]


        batch_size = source.shape[1]

        target_vocab_size = self.output_dim

        outputs = torch.zeros(target_len, batch_size, target_vocab_size).to(self.device)

        encoder_hidden = self.encoder(source, source_len)

        # mask = [batch_size, src len]
        # without sos token at the beginning and eos token at the end

        #x = target[0,:]
        decoder_input = torch.ones(batch_size).long().to(self.device)
        decoder_hidden = encoder_hidden
        encoder_outputs = None

        for t in range(target_len):
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
            outputs[t] = decoder_output

            #(N, english_vocab_size)
            #best_guess = output.argmax(1)
            topv, topi = decoder_output.topk(1)
            decoder_input = topi.squeeze().detach()

            decoder_input = target[t] if random.random() < teacher_force_ratio and target is not None else decoder_input

        return outputs

    def loss(self, logits, target):
        return self._loss(logits, target)

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=self.learning_rate)

        lr_scheduler = {
            'scheduler': optim.lr_scheduler.OneCycleLR(
                optimizer,
                max_lr = self.learning_rate,
                steps_per_epoch = 3379,
                epochs=self.max_epochs,
                anneal_strategy='linear',
                final_div_factor=1000,
                pct_start = 0.01
            ),
            "name": "learning_rate",
            "interval":"step",
            "frequency": 1
        }

        return [optimizer],[lr_scheduler]

    def training_step(self, batch, batch_idx):
        src_seq, trg_seq, src_lengths = batch['src'],batch['trg'], batch['src_len']

        output = self.forward(src_seq, src_lengths,trg_seq)

        # do not know if this is a problem, loss will be computed with sos token
        # without sos token at the beginning and eos token at the end

        output = output.view(-1, self.output_dim)


        trg_seq = trg_seq.transpose(0, 1)
        trg = trg_seq.reshape(-1)

        loss = self.loss(output, trg)

        self.log('train_loss',loss.item(),
                on_step = True,
                on_epoch=True,
                prog_bar = True,
                logger=True)

        return loss

    def validation_step(self, batch,batch_idx):
        """ validation is in eval model so we do not have to use placeholder input sensors"""
        src_seq, trg_seq, src_lengths = batch['src'],batch['trg'], batch['src_len']

        outputs = self.forward(src_seq, src_lengths, trg_seq, 0)

        logits = outputs[1:].view(-1, self.output_dim)

        trg = trg_seq[1:].reshape(-1)

        loss = self.loss(logits, trg)

        pred_seq = outputs[1:].argmax(2)   # seq_len*batch_size*vocab_size -> seq_len * batch_size

        # change layout: sesq_len * batch_size -> batch_size * seq_len
        pred_seq = pred_seq.T

        # change layout: seq_len * batch_size -> batch_size * seq_len
        trg_batch = trg_seq[1:].T

        # compare list of predicted ids for all sequences in a batch to targets
        acc = plfunc.accuracy(pred_seq.reshape(-1), trg_batch.reshape(-1))

        # need to cast to list of predicted sequences ( as list of token ids ) [ seq_tok1, seqtok2]
        predicted_ids - pred_seq.tolist()

        # need to add additional dim to each target reference sequence in order to
        # conver to format needed by blue_score_func
        # [seq1=[[reference1],[reference2]], seq2=[reference1]]
        target_ids = torch.unsqueeze(trg_batch, 1).tolist()

        bleu_score - plfunc.nlp.bleu_score(predicted_ids, target_ids, n_gram=3).to(self.device)

        self.log(
        'val_loss',
        loss,
        on_step=False,
        on_epoch=True,
        prog_bar=True,
        logger=True,
        sync_dist=True)

        self.log(
            "val_acc",
            acc,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            sync_dist=True
        )

        self.log(
            "val_bleu_idx",
            bleu_score,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            sync_dist=True
        )

        return loss, acc, bleu_score

[(20, 'torch.nn.Dropout', 'nn.Dropout', False, 'from torch import nn\n'), (21, 'torch.nn.Embedding', 'nn.Embedding', False, 'from torch import nn\n'), (22, 'torch.nn.LSTM', 'nn.LSTM', False, 'from torch import nn\n'), (51, 'torch.nn.Embedding', 'nn.Embedding', False, 'from torch import nn\n'), (52, 'torch.nn.Dropout', 'nn.Dropout', False, 'from torch import nn\n'), (53, 'torch.nn.LSTM', 'nn.LSTM', False, 'from torch import nn\n'), (54, 'torch.nn.Linear', 'nn.Linear', False, 'from torch import nn\n'), (127, 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', False, 'from torch import nn\n'), (210, 'torch.optim.lr_scheduler.OneCycleLR', 'optim.lr_scheduler.OneCycleLR', False, 'from torch import optim\n'), (150, 'torch.nn.init.normal_', 'nn.init.normal_', False, 'from torch import nn\n'), (153, 'torch.nn.init.constant_', 'nn.init.constant_', False, 'from torch import nn\n'), (178, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (279, 'torch.unsqueeze', 'torch.unsqueeze', False, 'import torch\n'), (281, 'pytorch_lightning.metrics.functional.nlp.bleu_score', 'plfunc.nlp.bleu_score', True, 'import pytorch_lightning.metrics.functional as plfunc\n'), (186, 'torch.ones', 'torch.ones', False, 'import torch\n'), (199, 'random.random', 'random.random', False, 'import random\n')]