#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    : 2019-08-15 18:31
# @Author  : apollo2mars
# @File    : BiDAF.py
# @Contact : apollo2mars@gmail.com
# @Desc    :

import tensorflow as tf
from tensorflow.python.ops import variable_scope as vs


class BiDAF(object):
    """
    Module for bidirectional attention flow.

    """

    def __init__(self, tokenizer):
        """
        Inputs:
          keep_prob: tensor containing a single scalar that is the keep probability (for dropout)
          vec_size: size of the word embeddings. int
        """
        self.tokenizer = tokenizer
        self.keep_prob = 1.0
        self.vec_size = len(tokenizer.word2idx)
        self.S_W = tf.get_variable('S_W', [self.vec_size*3], tf.float32,
            tf.contrib.layers.xavier_initializer())

    def build_graph(self, q, q_mask, c, c_mask):
        """
        Inputs:
          c: context matrix, shape: (batch_size, num_context_words, vec_size).
          c_mask: Tensor shape (batch_size, num_context_words).
            1s where there's real input, 0s where there's padding
          q: question matrix (batch_size, num_question_words, vec_size)
          q_mask: Tensor shape (batch_size, num_question_words).
            1s where there's real input, 0s where there's padding
          N = num_context_words
          M = Num_question_words
          vec_size = hidden_size * 2

        Outputs:
          output: Tensor shape (batch_size, N, vec_size*3).
            This is the attention output.
        """
        with vs.variable_scope("BiDAF"):

            # Calculating similarity matrix
            c_expand = tf.expand_dims(c,2)  #[batch,N,1,2h]
            print(c_expand)
            q_expand = tf.expand_dims(q,1)  #[batch,1,M,2h]
            print(q_expand)
            c_pointWise_q = c_expand * q_expand  #[batch,N,M,2h]
            print(c_pointWise_q)

            c_input = tf.tile(c_expand, [1, 1, tf.shape(q)[1], 1])
            q_input = tf.tile(q_expand, [1, tf.shape(c)[1], 1, 1])

            concat_input = tf.concat([c_input, q_input, c_pointWise_q], -1) # [batch,N,M,6h]
            print(concat_input)

            similarity = tf.reduce_sum(concat_input * self.S_W, axis=3)  #[batch,N,M]
            print(similarity)

            # Calculating context to question attention
            similarity_mask = tf.expand_dims(q_mask, 1) # shape (batch_size, 1, M)
            print(similarity_mask)
            _, c2q_dist = masked_softmax(similarity, similarity_mask, 2) # shape (batch_size, N, M). take softmax over q
            print(c2q_dist)

            # Use attention distribution to take weighted sum of values
            c2q = tf.matmul(c2q_dist, q) # shape (batch_size, N, vec_size)
            print(c2q)

            # Calculating question to context attention c_dash
            S_max = tf.reduce_max(similarity, axis=2) # shape (batch, N)
            print(S_max)
            _, c_dash_dist = masked_softmax(S_max, c_mask, 1) # distribution of shape (batch, N)
            print(c_dash_dist)
            c_dash_dist_expand = tf.expand_dims(c_dash_dist, 1) # shape (batch, 1, N)
            print(c_dash_dist_expand)
            c_dash = tf.matmul(c_dash_dist_expand, c) # shape (batch_size, 1, vec_size)
            print(c_dash)

            c_c2q = c * c2q # shape (batch, N, vec_size)
            print(c_c2q)

            c_c_dash = c * c_dash # shape (batch, N, vec_size)
            print(c_c_dash)

            # concatenate the output
            output = tf.concat([c2q, c_c2q, c_c_dash], axis=2) # (batch_size, N, vec_size * 3)
            print(output)


            # Apply dropout
            output = tf.nn.dropout(output, self.keep_prob)
            print(output)

            return output


def masked_softmax(logits, mask, dim):
    """
    Takes masked softmax over given dimension of logits.

    Inputs:
      logits: Numpy array. We want to take softmax over dimension dim.
      mask: Numpy array of same shape as logits.
        Has 1s where there's real data in logits, 0 where there's padding
      dim: int. dimension over which to take softmax

    Returns:
      masked_logits: Numpy array same shape as logits.
        This is the same as logits, but with 1e30 subtracted
        (i.e. very large negative number) in the padding locations.
      prob_dist: Numpy array same shape as logits.
        The result of taking softmax over masked_logits in given dimension.
        Should be 0 in padding locations.
        Should sum to 1 over given dimension.
    """
    exp_mask = (1 - tf.cast(mask, 'float')) * (-1e30) # -large where there's padding, 0 elsewhere
    masked_logits = tf.add(logits, exp_mask) # where there's padding, set logits to -large
    prob_dist = tf.nn.softmax(masked_logits, dim)
    return masked_logits, prob_dist