import torch
import torch.nn as nn
import torch.nn.functional as F
import random
from torch.autograd import Variable as Var

import numpy as np


# Embedding layer that has a lookup table of symbols that is [full_dict_size x input_dim]. Includes dropout.
# Works for both non-batched and batched inputs
class EmbeddingLayer(nn.Module):
    # Parameters: dimension of the word embeddings, number of words in vocabulary, and the dropout rate to apply
    # (0.2 is often a reasonable value)
    def __init__(self, input_dim, full_dict_size, embedding_dropout_rate):
        super(EmbeddingLayer, self).__init__()
        self.dropout = nn.Dropout(embedding_dropout_rate)
        self.word_embedding = nn.Embedding(full_dict_size, input_dim)

    # Takes either a non-batched input [sent len x input_dim] or a batched input
    # [batch size x sent len x input dim]
    def forward(self, input):
        embedded_words = self.word_embedding(input)
        final_embeddings = self.dropout(embedded_words)
        return final_embeddings


# One-layer RNN encoder for batched inputs -- handles multiple sentences at once. You're free to call it with a
# leading dimension of 1 (batch size 1) but it does expect this dimension.
class RNNEncoder(nn.Module):
    # Parameters: input size (should match embedding layer), hidden size for the LSTM, dropout rate for the RNN,
    # and a boolean flag for whether or not we're using a bidirectional encoder
    def __init__(self, input_size, hidden_size, dropout, bidirect):
        super(RNNEncoder, self).__init__()
        self.bidirect = bidirect
        self.input_size = input_size
        self.hidden_size = hidden_size
        #nn.Linear(hidden_size, output_size)
        self.reduce_h_W = nn.Linear(hidden_size * 2, hidden_size, bias=True)
        self.reduce_c_W = nn.Linear(hidden_size * 2, hidden_size, bias=True)
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True,
                               dropout=dropout, bidirectional=self.bidirect)
        self.init_weight()

    # Initializes weight matrices using Xavier initialization
    def init_weight(self):
        # ih: input - hidden
        # hh: hidden - hidden 
        # l0: layer 0
        # l1: layer 1
        # weight_... : weights of input_size x hidden_size or hidden_size x hidden_size (depending on whether ih or hh is specified) are initialized
        # bias_... : bias of size hidden_size initialized
        nn.init.xavier_uniform_(self.rnn.weight_hh_l0, gain=1)
        nn.init.xavier_uniform_(self.rnn.weight_ih_l0, gain=1)
        if self.bidirect:
            nn.init.xavier_uniform_(self.rnn.weight_hh_l0_reverse, gain=1)
            nn.init.xavier_uniform_(self.rnn.weight_ih_l0_reverse, gain=1)
        nn.init.constant_(self.rnn.bias_hh_l0, 0)
        nn.init.constant_(self.rnn.bias_ih_l0, 0)
        if self.bidirect:
            nn.init.constant_(self.rnn.bias_hh_l0_reverse, 0)
            nn.init.constant_(self.rnn.bias_ih_l0_reverse, 0)

    def get_output_size(self):
        return self.hidden_size * 2 if self.bidirect else self.hidden_size

    def sent_lens_to_mask(self, lens, max_length):
        return torch.from_numpy(np.asarray([[1 if j < lens.data[i].item() else 0 for j in range(0, max_length)] for i in range(0, lens.shape[0])]))

    # embedded_words should be a [batch size x sent len x input dim] tensor
    # input_lens is a tensor containing the length of each input sentence
    # Returns output (each word's representation), context_mask (a mask of 0s and 1s
    # reflecting where the model's output should be considered), and h_t, a *tuple* containing
    # the final states h and c from the encoder for each sentence.
    def forward(self, embedded_words, input_lens):
        # Takes the embedded sentences, "packs" them into an efficient Pytorch-internal representation
        packed_embedding = nn.utils.rnn.pack_padded_sequence(embedded_words, input_lens, batch_first=True)
        # Runs the RNN over each sequence. Returns output at each position as well as the last vectors of the RNN
        # state for each sentence (first/last vectors for bidirectional)
        output, hn = self.rnn(packed_embedding)
        # Unpacks the Pytorch representation into normal tensors
        output, sent_lens = nn.utils.rnn.pad_packed_sequence(output)
        # input_lens is a pytorch Variable which is a wrapping over a tensor that makes it differentiable. 
        # To unpack the underlying tensor, we use input_lens.data.
        max_length = input_lens.data[0].item()
        context_mask = self.sent_lens_to_mask(sent_lens, max_length)

        # Grabs the encoded representations out of hn, which is a weird tuple thing.
        # Note: if you want multiple LSTM layers, you'll need to change this to consult the penultimate layer
        # or gather representations from all layers.

        # output - (sentence_length,batch_size,hid_size*num_directions)
        # hn[0] - (num_layers*num_directions, batch_size, hid_size)
        # hn[1] - (num_layers*num_directions, batch_size, hid_size)

        if self.bidirect:
            h, c = hn[0], hn[1]
            # Grab the representations from forward and backward LSTMs
            h_, c_ = torch.cat((h[0], h[1]), dim=1), torch.cat((c[0], c[1]), dim=1)
            # Reduce them by multiplying by a weight matrix so that the hidden size sent to the decoder is the same
            # as the hidden size in the encoder
            new_h = self.reduce_h_W(h_)
            new_c = self.reduce_c_W(c_)
            h_t = (new_h, new_c)
        else:
            h, c = hn[0][0], hn[1][0] # hn[0][0] equivalent to hn[0].squeeze() for batch size = 1
            h_t = (h, c)
        return (output, context_mask, h_t)

class DecoderCell(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, dropout):
        super(DecoderCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True) 
        self.fc = nn.Linear(hidden_size, num_classes)
        self.init_weight()
    
    def init_weight(self):
        # ih: input - hidden, hh: hidden - hidden, l0: layer 0, l1: layer 1
        # weight_... : weights of input_size x hidden_size or hidden_size x hidden_size (depending on whether ih or hh is specified) are initialized
        # bias_... : bias of size hidden_size initialized
        nn.init.xavier_uniform_(self.rnn.weight_hh_l0, gain=1)
        nn.init.xavier_uniform_(self.rnn.weight_ih_l0, gain=1)
        nn.init.constant_(self.rnn.bias_hh_l0, 0)
        nn.init.constant_(self.rnn.bias_ih_l0, 0)

    def forward(self, embedded_input, input_lens, hidden_input, enc_output_states):
        packed_embedding = nn.utils.rnn.pack_padded_sequence(embedded_input, input_lens, batch_first=True)
        output, hn  = self.rnn(packed_embedding, hidden_input)        
        h, c = hn[0][0], hn[1][0] 
        h_t = (h, c)
        class_scores = self.fc(h)
        return (class_scores, h_t)

class AttentionDecoderCell(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, max_inp_len, dropout):
        super(AttentionDecoderCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
        self.partial_f = nn.Linear(hidden_size, hidden_size*2)
        self.to_class_scores = nn.Linear(hidden_size *3, num_classes, bias=True)
        self.init_weight()

    def init_weight(self):
        # ih: input - hidden, hh: hidden - hidden, l0: layer 0, l1: layer 1
        # weight_... : weights of input_size x hidden_size or hidden_size x hidden_size (depending on whether ih or hh is specified) are initialized
        # bias_... : bias of size hidden_size initialized
        nn.init.xavier_uniform_(self.rnn.weight_hh_l0, gain=1)
        nn.init.xavier_uniform_(self.rnn.weight_ih_l0, gain=1)
        nn.init.constant_(self.rnn.bias_hh_l0, 0)
        nn.init.constant_(self.rnn.bias_ih_l0, 0)

    def forward(self, embedded_input, input_lens, hidden_input, enc_output_states):
        packed_embedding = nn.utils.rnn.pack_padded_sequence(embedded_input, input_lens, batch_first=True)
        output, hn  = self.rnn(packed_embedding, hidden_input)        
        h, c = hn[0][0], hn[1][0] 
        h_t = (h, c)
        e_i = torch.mm(self.partial_f(h),enc_output_states.squeeze(1).t()) 
        a_i = torch.softmax(e_i,dim=1)
        c_i = torch.mm(a_i,enc_output_states.squeeze(1))
        c_i_final = torch.tanh(torch.cat((c_i,h), dim=1))
        class_scores = self.to_class_scores(c_i_final)
        return (class_scores, h_t)
