text: import os, sys, math
  sys.path.insert(0, '../../..')
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  import scipy.stats
  
  class Attention(nn.Module):
      def __init__(self, encoder_size, decoder_size, device, type="additive"):
          """ Attention module.         
                  TODO description for each type
                  
                  TODO needed bias=False for KVQ transformations, as well for type?
                  
                  TODO need mask?
                  
              Args:
                  encoder_size (int): Size of the encoder's output (as input for the decoder).
                  decoder_size (int): Size of the decoder's output.
                  device (torch.device): Device (eg. torch.device("cpu"))
                  type (string): One of several types of attention
          
              See: https://arxiv.org/pdf/1902.02181.pdf
              
              Notes: 
                  Self-Attention(Intra-attention) Relating different positions of the same input sequence. Theoretically the self-attention can adopt any score functions above, but just replace the target sequence with the same input sequence.
                  Global/Soft	Attending to the entire input state space.
                  Local/Hard	Attending to the part of input state space; i.e. a patch of the input image.
          """
          super(Attention, self).__init__()
          self.encoder_size = encoder_size
          self.decoder_size = decoder_size
          self.type = type
          
          # transforms encoder states into keys
          self.key_annotation_function = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
          # transforms encoder states into values 
          self.value_annotation_function = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
          # transforms the hidden state into query
          self.query_annotation_function = nn.Linear(self.decoder_size, self.encoder_size, bias=False) # NOTE: transforming q to K size 
          
          if type == "additive":
              # f(q, K) = wimp*tanh(W1K + W2q + b) , Bahdanau et al., 2015            
              self.V = nn.Linear(self.encoder_size, 1, bias=False) 
              self.W1 = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
              self.W2 = nn.Linear(self.encoder_size, self.encoder_size, bias=False) # encoder size because q is now K's size, otherwise dec_size to enc_size
              self.b = nn.Parameter(torch.zeros(self.encoder_size))
              
          elif type == "coverage": # https://arxiv.org/pdf/1601.04811.pdf
              # f(q, K) = wimp*tanh(W1K + W2q + b) , Bahdanau et al., 2015            
              self.V = nn.Linear(self.encoder_size, 1, bias=False) 
              self.W1 = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
              self.W2 = nn.Linear(self.encoder_size, self.encoder_size, bias=False) # encoder size because q is now K's size, otherwise dec_size to enc_size
              self.b = nn.Parameter(torch.zeros(self.encoder_size))
              
              self.coverage_dim = 10
              self.coverage_input_size = self.coverage_dim + 1 + self.encoder_size + self.encoder_size
              self.cov_gru = nn.GRU(self.coverage_input_size, self.coverage_dim, batch_first=True)
              self.W3 = nn.Linear(self.coverage_dim, self.encoder_size)
                                 
          elif (type == "multiplicative" or type == "dot"):
              # f(q, K) = q^t K , Luong et al., 2015
              # direct dot product, nothing to declare here
              pass
              
          elif type == "scaled multiplicative" or type == "scaled dot":
              # f(q, K) = multiplicative / sqrt(dk) , Vaswani et al., 2017
              self.scale = math.sqrt(self.encoder_size)            
              
          elif type == "general" or type == "bilinear":
              # f(q, K) = q^t WK , Luong et al., 2015
              self.W = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
              
          elif type == "biased general":
              # f(q, K) = K|(W q + b) Sordoni et al., 2016
              pass
          elif type == "activated general":
              # f(q, K) = act(q|WK + b) Ma et al., 2017        
              pass
          elif type == "concat":
              # f(q, K) = act(W[K;q] + b) , Luong et al., 2015
              pass
          elif type == "p":
              # https://arxiv.org/pdf/1702.04521.pdf pagina 3, de adaugat predict-ul in attention, KVP si Q
              pass
          else:
              raise Exception("Attention type not properly defined! (got type={})".format(self.type))
          self.device = device
          self.to(self.device)
  
      def _reshape_state_h(self, state_h):    
          """
          Reshapes the hidden state to desired shape
          Input: [num_layers * 1, batch_size, decoder_hidden_size]  
          Output: [batch_size, 1, decoder_hidden_state]
  
          Args:
              state_h (tensor): Hidden state of the decoder.        
                  [num_layers * 1, batch_size, decoder_hidden_size]
          
          Returns:
              The reshaped hidden state.
                  [batch_size, 1, decoder_hidden_state]
          """
          num_layers, batch_size, hidden_size = state_h.size()
          # in case the decoder has more than 1 layer, take only the last one -> [1, batch_size, decoder_hidden_size]
          if num_layers > 1:
              state_h = state_h[num_layers-1:num_layers,:,:]
  
          # [1, batch_size, decoder_hidden_size] -> [batch_size, 1, decoder_hidden_size]
          return state_h.permute(1, 0, 2)
  
      def reset_coverage(self, batch_size, enc_seq_len, force_weight): # called on every new batch
          self.force_weight = force_weight
          
          if self.type == "coverage":
              self.C = torch.zeros(batch_size, enc_seq_len, self.coverage_dim, device=self.device)
              self.gru_input = torch.zeros(batch_size, 1, self.coverage_input_size, device=self.device)
                      
      def _coverage_compute_next_C(self, attention_weights, enc_output, state_h):
          # attention_weights:        [batch_size, seq_len, 1]
          # enc_output :              [batch_size, seq_len, encoder_size]
          # state_h (after reshape):  [batch_size, 1, encoder_size]       
          # self.C at prev timestep:  [batch_size, seq_len, coverage_dim]
          seq_len = enc_output.size[1]
          for i in range(seq_len):
              self.gru_input[:,:,0:self.coverage_dim] = self.C[:,i:i+1,:] # cat self.C at prev timestep for position i of source word
              self.gru_input[:,:,self.coverage_dim:self.coverage_dim+1] = attention_weights[:,i:i+1,:] # cat attention weight
              self.gru_input[:,:,self.coverage_dim+1:self.coverage_dim+1+self.encoder_size] = enc_output[:,i:i+1,:] # cat encoder output
              self.gru_input[:,:,self.coverage_dim+1+self.encoder_size:] = state_h[:,0:1,:] # cat state_h
              self.C[:,i:i+1,:] = self.cov_gru(self.gru_input)
          
      def _energy (self, K, Q):
          """ 
              Calculates the compatibility function f(query, keys)
              
              Args:
                  K (tensor): Keys tensor of size [batch_size, seq_len, encoder_size]
                  Q (tensor): Query tensor of size [batch_size, 1, decoder_size], but now dec_size is enc_size due to Q annotation
                  
              Returns: 
                  energy tensor of size [batch_size, seq_len, 1]
          """
          if self.type == "additive":                        
              return self.V(torch.tanh(self.W1(K) + self.W2(Q) + self.b))
          
          elif self.type == "coverage":
              return self.V(torch.tanh(self.W1(K) + self.W2(Q) + self.W3(self.C) + self.b))
              
          elif self.type == "multiplicative" or self.type == "dot":    
              # q^t K means batch matrix multiplying K with q transposed:
              # bmm( [batch_size, seq_len, enc_size] , [batch_size, enc_size, 1] ) -> [batch_size, seq_len, 1]            
              return torch.bmm(K, Q.transpose(1,2))
          
          elif self.type == "scaled multiplicative" or self.type == "scaled dot":
              # same as multiplicative but scaled
              return torch.bmm(K, Q.transpose(1,2))/self.scale
      
          elif self.type == "general" or self.type == "bilinear":
              # f(q, K) = q^t WK , Luong et al., 2015
              return torch.bmm(self.W(K), Q.transpose(1,2))
  
      def forward(self, enc_output, state_h, decoder_step, dec_seq_len, mask=None):
          """
          This function calculates the context vector of the attention layer, given the hidden state and the encoder
          last lstm layer output.
  
          Args:
              state_h (tensor): The raw hidden state of the decoder's LSTM 
                  Shape: [num_layers * 1, batch_size, decoder_size].
              enc_output (tensor): The output of the last LSTM encoder layer. 
                  Shape: [batch_size, seq_len, encoder_size].
              mask (tensor): 1 and 0 as for encoder input
                  Shape: [batch_size, seq_len].
  
          Returns:
              context (tensor): The context vector. Shape: [batch_size, encoder_size]
              attention_weights (tensor): Attention weights. Shape: [batch_size, seq_len, 1]
          """
          batch_size = enc_output.shape[0]
          seq_len = enc_output.shape[1]        
          state_h = self._reshape_state_h(state_h) # [batch_size, 1, decoder_size]
          
          # get K, V, Q
          K = self.key_annotation_function(enc_output) # [batch_size, seq_len, encoder_size]
          V = self.value_annotation_function(enc_output) # [batch_size, seq_len, encoder_size]
          Q = self.query_annotation_function(state_h) # [batch_size, 1, encoder_size]
          
          # calculate energy
          energy = self._energy(K,Q) # [batch_size, seq_len, 1]        
          
          # mask with -inf paddings
          if mask is not None:            
              energy.masked_fill_(mask.unsqueeze(-1) == 0, -np.inf)
          
          # transform energy into probability distribution using softmax        
          attention_weights = torch.softmax(energy, dim=1) # [batch_size, seq_len, 1]
          
          # calculate gaussian distribution as forced probability, assume that decoder len is roughly the same as encoder len        
          """
          if self.force_weight > 0.:
              approx_position = decoder_step*seq_len/dec_seq_len
              x = np.linspace(0, seq_len, seq_len)
              y = scipy.stats.norm.pdf(x, approx_position, 3) # loc (mean) is decoder_step, scale (std dev) = 1.        
              #print(decoder_step)
              #print(np.sum(y))
              y = y / np.sum(y) # rescale to make it a PDF
              gaussian_dist = torch.tensor(y, dtype = attention_weights.dtype, device = self.device) # make it a tensor, it's [seq_len]
              gaussian_dist = gaussian_dist.repeat(batch_size, 1) # same for all examples in batch, now it's [batch_size, seq_len]
          
              # calculate attention_weights as a PDF 
              attention_weights = self.force_weight * gaussian_dist + (1-self.force_weight) * attention_weights.squeeze(2)
              attention_weights = attention_weights.unsqueeze(2) # recreate [batch_size, seq_len, 1]
          """
          
          # for coverage only, calculate the next C
          if type=="coverage":
              self._coverage_compute_next_C(attention_weights, enc_output, state_h)            
          
          # calculate weighted values z (element wise multiplication of energy * values)        
          # attention_weights is [batch_size, seq_len, 1], V is [batch_size, seq_len, encoder_size], z is same as V
          z = attention_weights*V # same as torch.mul(), element wise multiplication
          
          # finally, calculate context as the esum of z. 
          # z is [batch_size, seq_len, encoder_size], context will be [batch_size, encoder_size]
          context = torch.sum(z, dim=1)
          
          return context, attention_weights # [batch_size, encoder_size], [batch_size, seq_len, 1]
  
  
  if __name__ == "__main__":
      import numpy as np
      
      
      # debug stuff:    
      q = torch.tensor([ [ [2.,2.,2.] ] ])
      K = torch.tensor([ [ [1.,1.,1.] , [5.,5.,5.] ] ])
      #result would be ([ [ [2.,2.,2.] , [2.5,2.5,2.5] ] ])
      
      print(K.size())
      print(q)
      print(q.size())
      #qt = q.expand(1,3,3)#q.transpose(1,2)
      qt = q.transpose(1,2)
      print(qt)
      print(qt.size())    
      
      print()
      r = torch.bmm(K,qt)
      print(r)
      print(r.size())
      print()
      #print(e1.size())
      #print(v1.size())
      #qq = e1*v1
      #print(qq)
      
      
      
      # prep inputs
      batch_size = 2
      seq_len = 10
      enc_size = 4
      dec_layers = 5
      dec_size = 3
      
      encoder_outputs = torch.tensor(np.random.rand(batch_size, seq_len, enc_size), dtype=torch.float)
      decoder_hidden_state = torch.tensor(np.random.rand(dec_layers*1, batch_size, dec_size), dtype=torch.float) # 1 for unidirectional
      
      # prep layer
      device = torch.device("cpu")
      #type = "additive"    
      type = "general"    
      att = Attention(enc_size, dec_size, device, type)
      
      # run
      context, attention_weights = att(encoder_outputs, decoder_hidden_state)
      print("Output is:")
      print(context)
      print("Attention weights size:" + str(attention_weights.size()))
      
      # debug stuff:    
      #e1 = torch.tensor([[[2],[0.5]]])
      #v1 = torch.tensor([ [ [1.,1.,1.] , [5.,5.,5.] ] ])
      #result would be ([ [ [2.,2.,2.] , [2.5,2.5,2.5] ] ])
      #print(e1.size())
      #print(v1.size())
      #qq = e1*v1
      #print(qq)
  
      