# -*- coding: utf-8 -*-
"""pointer_gen_modules.ipynb

Automatically generated by Colaboratory.

"""

import numpy as np
import random
import tensorflow as tf
import tensorflow.nn as nn

from utils import Linear
from utils import apply_mask_normalize
from utils import _mask_and_avg
from utils import _calc_final_dist

class Encoder():
  """ A simple encoder class to encode the input via Bidirectional LSTM
      Args:
          hpm : hyperparameters
          rand_unif_init : Initializer Object (random uniform) to initialize LSTMs parameters
          rand_norm_init : Initializer object (truncate normal) to initialize weights and biases for linear transf. 
  """
  def __init__(self, hpm, rand_unif_init, rand_norm_init):
    self.hpm= hpm
    self.rand_unif_init = rand_unif_init
    self.rand_norm_init = rand_norm_init
    
    with tf.variable_scope('encoder'):
      self.lstm_cell_fw = tf.contrib.rnn.LSTMCell(self.hpm["hidden_size"],
                                           state_is_tuple= True, initializer=self.rand_unif_init) # forward lstm cell
      self.lstm_cell_bw = tf.contrib.rnn.LSTMCell(self.hpm["hidden_size"],
                                           state_is_tuple= True, initializer=self.rand_unif_init) # backward lstm cell
      
      self.w_c = Linear(self.hpm['hidden_size'], True, "reduce_c", self.rand_norm_init) # Parameters for the concatenated state linear transf.
      self.w_h = Linear(self.hpm['hidden_size'], True, 'reduce_h', self.rand_norm_init) # Parameters for the concatenated hidden output linear transf.
      
    
  
  def __call__(self, encoder_inputs, seq_lens):
    """ Call method for the encoding feedforward 
        Args:
            encoder_inpputs : 3D tensor, shape : [batch_size, max_enc_len, embed_size]
            seq_lens : 1D tensor, lengths of the sequences (without padding) in the batch, shape : [batch_size]
            
        Returns:
            encoder_outputs : 3D tensor, output of the bidirectional dynamic rnn, shape : [batch_size, None, 2*hidden_size] (None because the max seq len vary with the batch)
            new state : tuple object made of two tensors : c => state, h=> last hidden output, shape : [2,batch_size, hidden_size]
    """
    with tf.variable_scope('encoder', reuse = tf.AUTO_REUSE):
      (encoder_outputs, (fw_st,bw_st)) = tf.nn.bidirectional_dynamic_rnn(
                                      self.lstm_cell_fw, self.lstm_cell_bw, encoder_inputs, 
                                      dtype=tf.float32, swap_memory=True, 
                                    sequence_length = seq_lens) 
    
      encoder_outputs=tf.concat(encoder_outputs, axis= 2)
    
      old_c= tf.concat(values=[fw_st.c,bw_st.c], axis= 1) # we concatenate the forward and backward state, shape: [batch_size, 2*hidden_size]
      old_h= tf.concat(values=[fw_st.h,bw_st.h], axis= 1) # we concatenate the forwarrd and backward last hidden output, shape : [batch_size, 2*hidden_size]
      new_c= tf.nn.relu(self.w_c(old_c)) # linear transformation + relu activation, shape : [batch_size, hidden_size]
      new_h= tf.nn.relu(self.w_h(old_h)) # same as above
    
    return encoder_outputs, tf.contrib.rnn.LSTMStateTuple(new_c,new_h)



class Decoder():
  """
    A simple decoder class made of a unidirectional LSTM cell which decodes the next word given a previous one, a context vector and a previous state
    Args : 
        hpm : hyperparameters
        rand_unif_init : Initializer Object (random uniform) to initialize LSTM parameters
  """
  def __init__(self,hpm,rand_unif_init):
    self.hpm= hpm
    self.rand_unif_init = rand_unif_init

    with tf.variable_scope('decoder'):
      self.lstm_cell= tf.contrib.rnn.LSTMCell(self.hpm["hidden_size"],
                                           state_is_tuple= True, initializer=self.rand_unif_init) # unidirectional lstm cell
  

  def __call__(self, dec_inputs, prev_state):
    """ Feedforward method for the simple decoder
    
        Args:
            dec_inputs : 2D tensor, list of words time step t for each sequence in the batch, shape = [batch_size, embed_size]
            prev_state : tuple object made of two vectors : c => state, h => last hidden output, shape : [2, batch_size, hidden_size]
            
        Returns:
            decoder_outputs : 2D tensor, shape = [batch_size, hidden_size]
            curr_st : current state of the decoder, shape : [2, batch_size, hidden_size]
    """
    with tf.variable_scope('decoder', reuse = tf.AUTO_REUSE):
    	decoder_outputs, curr_st= tf.nn.dynamic_rnn(self.lstm_cell, dec_inputs,
                                           dtype= tf.float32, initial_state=prev_state, swap_memory= True, time_major=True)
    return decoder_outputs, curr_st



class Attention_decoder():
  """
      An attentional based encoder-decoder model (bhadanau attention, additive style)
      Args:
          hpm : hyperparameters
          rand_unif_init : Initializer Object (random uniform) to initialize LSTMs parameters
          rand_norm_init : Initializer object (truncate normal) to initialize weights and biases for linear transf. 
          
  """
  def __init__(self,hpm, rand_unif_init, rand_norm_init ):
    self.rand_unif_init = rand_unif_init
    self.rand_norm_init = rand_norm_init
    self.hpm=hpm
    
    with tf.variable_scope('attention_decoder', reuse = tf.AUTO_REUSE):
      self.decoder= Decoder(self.hpm, self.rand_unif_init) # simple decoder object (unidirecitional lstm)
    
      # Almost all the parameters (weights and biases) for the linear transformations (see below in the call method)
    
      self.w_h = Linear(self.hpm['attn_hidden_size'], True, "h")
      self.w_s = Linear(self.hpm['attn_hidden_size'], True, "s" )
      self.v = Linear(1, False, 'V')
    
      self.w_dec = Linear(self.hpm['emb_size'],True, "dec_inp")
      self.w_out = Linear(self.hpm['vocab_size'], True, 'out')
 
      if self.hpm['pointer_gen']:
        self.w_c_reduce = Linear(1, True, 'c_reduce')
        self.w_s_reduce = Linear(1, True, 's_reduce')
        self.w_i_reduce = Linear(1, True, 'i_reduce')
   
 

  def __call__(self, enc_outputs, enc_mask, enc_state, decoder_inputs,batch_max_oov_len = None, encoder_input_with_oov = None, cov_vec=None):
    """
        Attentional feedforward graph .
        We call this method once during training for each batch, and max_dec_len times for decode mode.
        
        Args:
            enc_outputs : 3D tensor, encoder outputs, shape : [batch_size, batch_max_enc_len, 2*hidden_size]
            enc_mask : 2D tensor, encoder sequence mask, shape : [batch_size, batch_max_enc_len]
            decoder_inputs: 3D tensor, decoder inputs, shape : [batch_size, max_dec_len, embed_size]
            batch_max_oov_len : Integer, Maximum number of oov for the current batch, (None if pointer_gen = False)
            encoder_input_with_oov : 2D tensor, encoder input with oovs ids, shape : [batch_size, batch_max_enc_len]
            
            !!! NB : batch_max_enc_len is None when we build graph, and vary during the feedforward with the current batch treated, 
                      it is the maximum length of sequences of the current batch
                      
        Returns : A dictionary
            output : list max_dec_en of 2D tensors of shape [batch_size, vocab_size + batch_max_oov_len (if pointer_gen)]
            last_context_vector : 2D tensor, shape : [batch_size, 2*hidden_size], this will be useful in the decode mode
            dec_state : 2D tensor, decoder last state, shape : [2, batch_size, hidden_size]
            p_gen : max_dec_len-many list of 1D tensors of length[batch_size] (only if pointer_gen is true)
            attention_vec : max_dec_len-many list of 2D tensors of shape [batch_size, batch_max_enc_len] (only if coverage is true)
    """

    if(self.hpm["pointer_gen"]):
      p_gens=[] # if pointer gen, we add an array to store the probability of each word in the sequences to be generated or pointed on
     
    attn_dists = [] # array to store the attention distributions over the enc seq 
    dec_state = enc_state # we init the decoder state with the encoder last state
    outputs=[] # array to store the final probability distributions (decoded sequence)
    dec_inp = tf.unstack(decoder_inputs) # we unstack the decoder input to be able to enumerate over this tensor
    
    if  self.hpm['decode_using_prev']:
      argmax_arr = []
      samples_arr = []
      argmax_logprob_arr = []
      samples_logprob_arr = []

    # nested function
    def attention(dec_state, cov_vec=None):
      """
          Attention mechanism
          
          Args:
              dec_state : previous state of the decoder. shape : [2, batch_size, hidden_size]. For the first step, it corresponds to the encoder last state
              cov_vec : only if coverage is True (default None).  shape : [batch_size, <batch_max_enc_len>]. The previous coverage vector.
              
          Returns:
              attn_vec : 2D tensor, the attention vector at time step t. shape : [batch_size, <batch_max_enc_len>]
              context_vector : 2D tensor, shape: [batch_size, 2*hidden_size]
              cov_vec : 2D tensor, shape : [batch_size, <batch_max_enc_len>], the current coverage vector
      """
      if(self.hpm["coverage"]):
        with tf.variable_scope('coverage', reuse = tf.AUTO_REUSE ):
          w_c = tf.get_variable("w_c", [1,1,1,self.hpm['attn_hidden_size']]) # we add additional parameters for the coverage vector linear transf.
          
        cov_features = tf.expand_dims(tf.expand_dims(cov_vec, axis=2),axis=2) # given that the encoder max length is unknown and variable, we cannot just apply a 
        cov_features = tf.nn.conv2d(cov_features, w_c, [1,1,1,1], "SAME")     # linear transformation as above. To avoid this issue, we can apply a convolution layer
                                                                              # which will transform the cov vector as a simple linear transf. would.
        
      # e = V*tanh(w_h*h + w_s*s + w_c*c ) (the last term, only is coverage = True)
      # attention weights all over the encoder input sequence
      # shape : [batch_size, <batch_max_enc_len>, 1]
        e=tf.nn.tanh(self.w_h(enc_outputs) + 
                   tf.expand_dims(self.w_s(dec_state.c), axis=1) +
                   tf.squeeze(cov_features, [2]))
      else:
        e=tf.nn.tanh(self.w_h(enc_outputs) + 
                   tf.expand_dims(self.w_s(dec_state.c), axis=1))
      e = self.v(e)
      
      # we take off the last dimension which equals 1 
      e =  tf.reshape(e, [ e.get_shape().as_list()[0], -1]) # shape : [batch_size, <batch_max_enc_len>]

      
      attn_vec = tf.nn.softmax(e, axis=-1) # we apply a softmax on the attention weights to normalize them and obtain the attention vector.
      attn_vec = apply_mask_normalize(attn_vec, enc_mask) # Given that the input is padded with <PAD> token, the attentions weights over those tokens
                                                          # are not relevant, we apply the encoder input masks on the attention vectors to drop those 'irrelevant' attention weights
                                                          # and finally we re-normalize the attention weights to obtain probability distributions
      
      # context vector computation
      # we multiply the encoder outputs by the attention vector weigths (a weight for each output vector, when we consider only one sequence for the example)
      weighted_enc_outputs = tf.multiply(enc_outputs, tf.expand_dims(attn_vec, axis=-1)) # context vector at time step t, shape : [batch_size, ]
      context_vec = tf.reduce_sum(weighted_enc_outputs, axis=1)
      
      if self.hpm['coverage']:
          cov_vec = cov_vec + attn_vec # we update the coverage
      
      return attn_vec, context_vec, cov_vec
      # end of nested function
      
    with tf.variable_scope('attention_decoder', reuse = tf.AUTO_REUSE):
      # we compute the initial context vector
      _ , context_vec, _  = attention( dec_state, cov_vec)
      timesteps = self.hpm['max_dec_len']
      decoder_input = dec_inp[0]
      a=0
      if not self.hpm['decode_using_prev']:
        a = 1
      for i in range (a,timesteps):
        # for each item in the decoder inputs (this loops only once for decode mode)
          
        # concatenation of input (previous word) and context vector at timestep t 
        new_dec_inp = tf.concat([decoder_input, context_vec], axis = -1) # shape : [batch_size, embed_size+2*hidden_size]
        new_dec_inp = self.w_dec(new_dec_inp) #shape : [batch_size, embed_size]

        # We apply the LSTM decoder on the new input
        dec_output, dec_state = self.decoder(tf.expand_dims(new_dec_inp, axis=0), dec_state) # dec_output shape : [1, batch_size, hidden_size]
                                                                                           # dec_state shape : [2, batch_size, hidden_size] (2 for the state c and the last hidden output h)
        # attention vector of the current step, context vector for the next step
        # we update the coverage vector
        attn_vec, context_vec, cov_vec  = attention( dec_state, cov_vec)
        attn_dists.append(attn_vec)
      
        dec_output = tf.reshape(dec_output, [-1, dec_output.get_shape().as_list()[-1]]) # shape : [batch_size, hidden_size]
        dec_output = self.w_out(dec_output) # shape : [batch_size, vocab_size]
        vocab_dist = dec_output
      
        if not self.hpm['pointer_gen']:
          outputs.append(vocab_dist) # we do not apply yet the softmax function because this function is integrated in some futures ops like the loss function
        else:
          # if pointer_gen=True, we need to compute the softmax function because of the scatter op with the attention distribution
          outputs.append(tf.nn.softmax(dec_output, axis=-1))
          state = tf.concat([dec_state.c, dec_state.h], axis=1)
        
          #p_gen computation with the current concatenated state, context vector and the decoder input
          p_gen = tf.nn.sigmoid(self.w_c_reduce(context_vec)+
                             self.w_s_reduce(state )+
                             self.w_i_reduce(new_dec_inp)) # shape : [batch_size, 1]
          p_gens.append(p_gen)
        
        

      if self.hpm['pointer_gen']:
        # we apply the scatter op between the output distibutions (over the vocabulary) with the attention distributions
        outputs = _calc_final_dist(encoder_input_with_oov, outputs, attn_dists, p_gens, batch_max_oov_len, self.hpm)

      
      if not self.hpm['decode_using_prev']:
        decoder_input = dec_inp[i]
      else:

        batch_nums = tf.range(0, limit=self.hpm['batch_size'], dtype=tf.int64)
        argmax_seqs = []
        argmax_seqs_log_probs = []
        for i , x in enumerate(outputs):
          max_ids = tf.argmax(x, axis=-1)
          indices = tf.stack((batch_nums, max_ids), axis = -1)
          log_probs = tf.gather_nd(x, indices)
          argmax_seqs.append(max_ids)
          argmax_seqs_log_probs.append(log_probs)


        soft_outputs = tf.stack(outputs)
        if not self.hpm['pointer_gen']:
          soft_outputs = tf.softmax(soft_outputs)

        argmax_seqs = tf.stack(argmax_seqs)
        argmax_seqs_log_probs = tf.stack(argmax_seqs_log_probs)
        
        sampler = tf.distributions.Categorical(logits=soft_outputs)
        samples = sampler.sample()
        samples_log_probs = sampler.log_prob(samples)
        samples_log_probs = tf.identity(samples_log_probs)

        argmax_arr.append(argmax_seqs)
        argmax_logprob_arr.append(argmax_seqs_log_probs)
        samples_arr.append(samples)
        samples_logprob_arr.append(samples_log_probs)

        decoder_input = samples

    if  self.hpm['decode_using_prev']:
      argmax_arr = tf.stack(argmax_arr)
      argmax_logprob_arr = tf.stack(argmax_logprob_arr)
      samples_arr = tf.stack(samples_arr)
      samples_logprob_arr = tf.stack(samples_logprob_arr)

    dic = { 'output':outputs, 'last_context_vector':context_vec, 'dec_state':dec_state, 'attention_vec':attn_dists} 
    if(self.hpm['pointer_gen']):
      dic['p_gen'] = p_gens
    if(self.hpm['coverage']):
      dic['coverage'] = cov_vec

    if  self.hpm['decode_using_prev']:
      dic.update({
        "argmax_seqs" : argmax_arr,
        "argmax_log_probs" : argmax_logprob_arr,
        "samples_seqs" : samples_arr,
        "samples_log_probs" : samples_logprob_arr
        })

    return dic