import os
from torch import nn
import torch.utils.model_zoo as model_zoo


model_urls = {
    'wmt-lstm' : 'https://s3.amazonaws.com/research.metamind.io/cove/wmtlstm-b142a7f2.pth'
}

model_cache =os.path.join(os.path.dirname(os.path.realpath(__file__)), '.torch')


class MTLSTM(nn.Module):

    def __init__(self, n_vocab=None, vectors=None, residual_embeddings=False):
        """Initialize an MTLSTM.
         
        Arguments:
            n_vocab (bool): If not None, initialize MTLSTM with an embedding matrix with n_vocab vectors
            vectors (Float Tensor): If not None, initialize embedding matrix with specified vectors
            residual_embedding (bool): If True, concatenate the input embeddings with MTLSTM outputs during forward
        """
        super(MTLSTM, self).__init__()
        self.embed = False
        if n_vocab is not None:
            self.embed = True
            self.vectors = nn.Embedding(n_vocab, 300)
            if vectors is not None:
                self.vectors.weight.data = vectors
        self.rnn = nn.LSTM(300, 300, num_layers=2, bidirectional=True, batch_first=True)
        self.rnn.load_state_dict(model_zoo.load_url(model_urls['wmt-lstm'], model_dir=model_cache))
        self.residual_embeddings = residual_embeddings

    def forward(self, inputs, lengths=None, hidden=None):
        """A pretrained MT-LSTM (McCann et. al. 2017). 
        This LSTM was trained with 300d 840B GloVe on the WMT 2017 machine translation dataset.
     
        Arguments:
            inputs (Tensor): If MTLSTM handles embedding, a Long Tensor of size (batch_size, timesteps).
                             Otherwise, a Float Tensor of size (batch_size, timesteps, features).
            lengths (Long Tensor): (batch_size, lengths) lenghts of each sequence for handling padding
            hidden (Float Tensor): initial hidden state of the LSTM
        """
        outputs, hidden_t = self.rnn(inputs, hidden)
        
        return outputs