import numpy as np
import tensorflow as tf
H = 2
N = 2
M = 3
BS = 10

def my_softmax(arr):
    max_elements = np.reshape(np.max(arr, axis = 2), (BS, N, 1))
    arr = arr - max_elements
    exp_array = np.exp(arr)
    print (exp_array)
    sum_array = np.reshape(np.sum(exp_array, axis=2), (BS, N, 1))
    return exp_array /sum_array

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, 'float64')) * (-1e30) # -large where there's padding, 0 elsewhere
    print (exp_mask)
    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


def test_build_similarity(contexts, questions):
    w_sim_1 = tf.get_variable('w_sim_1',
        initializer=w_1) # 2 * H
    w_sim_2 = tf.get_variable('w_sim_2',
        initializer=w_2) # 2 * self.hidden_size
    w_sim_3 = tf.get_variable('w_sim_3',
        initializer=w_3) # 2 * self.hidden_size
    q_tile = tf.tile(tf.expand_dims(questions, 0), [N, 1, 1, 1]) #  N x BS x M x 2H
    q_tile = tf.transpose(q_tile, (1, 0, 3, 2)) # BS x N x 2H x M
    contexts = tf.expand_dims(contexts, -1) # BS x N x 2H x 1
    result = (contexts * q_tile) # BS x N x 2H x M
    tf.assert_equal(tf.shape(result), [BS, N, 2 * H, M])
    result = tf.transpose(result, (0, 1, 3, 2)) # BS x N x M x 2H
    result = tf.reshape(result, (-1, N * M, 2 * H)) # BS x (NxM) x 2H
    tf.assert_equal(tf.shape(result), [BS, N*M, 2*H])

    # w_sim_1 = tf.tile(tf.expand_dims(w_sim_1, 0), [BS, 1])
    # w_sim_2 = tf.tile(tf.expand_dims(w_sim_2, 0), [BS, 1])
    # w_sim_3 = tf.tile(tf.expand_dims(w_sim_3, 0), [BS, 1])
    term1 = tf.matmul(tf.reshape(contexts, (BS * N, 2*H)), tf.expand_dims(w_sim_1, -1)) # BS x N
    term1 = tf.reshape(term1, (-1, N))
    term2 = tf.matmul(tf.reshape(questions, (BS * M, 2*H)), tf.expand_dims(w_sim_2, -1)) # BS x M
    term2 = tf.reshape(term2, (-1, M))
    term3 = tf.matmul(tf.reshape(result, (BS * N * M, 2* H)), tf.expand_dims(w_sim_3, -1))
    term3 = tf.reshape(term3, (-1, N, M)) # BS x N x M
    S = tf.reshape(term1,(-1, N, 1)) + term3 + tf.reshape(term2, (-1, 1, M))
    return S

def test_build_sim_mask():
    context_mask = np.array([True, True]) # BS x N
    question_mask = np.array([True, True, False]) # BS x M
    context_mask = np.tile(context_mask, [BS, 1])
    question_mask = np.tile(question_mask, [BS, 1])
    context_mask = tf.get_variable('context_mask', initializer=context_mask)
    question_mask = tf.get_variable('question_mask', initializer=question_mask)
    context_mask = tf.expand_dims(context_mask, -1) # BS x N x 1
    question_mask = tf.expand_dims(question_mask, -1) # BS x M x 1
    question_mask = tf.transpose(question_mask, (0, 2, 1)) # BS x 1 x M
    sim_mask = tf.matmul(tf.cast(context_mask, dtype=tf.int32),
            tf.cast(question_mask, dtype=tf.int32)) # BS x N x M
    return sim_mask

def test_build_c2q(S, S_mask, questions):
    _, alpha = masked_softmax(S, mask, 2) # BS x N x M
    return tf.matmul(alpha, questions)

def test_build_q2c(S, S_mask, contexts):
    # S = BS x N x M
    # contexts = BS x N x 2H
    m = tf.reduce_max(S * tf.cast(S_mask, dtype=tf.float64), axis=2) # BS x N
    beta = tf.expand_dims(tf.nn.softmax(m), -1) # BS x N x 1
    beta = tf.transpose(beta, (0, 2, 1))
    q2c = tf.matmul(beta, contexts)
    return m, beta, q2c

def test_concatenation(c2q, q2c):
    q2c = tf.tile(q2c, (1, N, 1))
    output = tf.concat([c2q, q2c], axis=2)
    tf.assert_equal(tf.shape(output), [BS, N, 4*H])
    return output

if __name__== "__main__":
    w_1 = np.array([1., 2., 3., 4.])
    w_2 = np.array([5., 6., 7., 8.])
    w_3 = np.array([13., 12., 11., 10.])

    c = np.array([[[1., 2., 3., 4.], [5., 6., 7., 8.]]]) # BS x N x 2H
    q = np.array([[[1., 2., 3., 0.], [5., 6., 7., 4.], [8., 9. , 10., 11.]]]) # BS x M x 2H
    c = np.tile(c, [BS, 1, 1])
    q = np.tile(q, [BS, 1, 1])


    questions = tf.get_variable('questions', initializer=q)
    contexts = tf.get_variable('contexts', initializer=c)

    S = test_build_similarity(contexts, questions)
    mask = test_build_sim_mask()
    c2q = test_build_c2q(S, mask, questions)
    m, beta, q2c = test_build_q2c(S, mask, contexts)
    output = test_concatenation(c2q, q2c)

    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        S_result, mask_result, c2q_r = sess.run([S, mask, c2q])
        actual_result = np.tile(np.array([[228, 772, 1372], [548, 1828, 3140]]), [BS, 1, 1])
        assert np.array_equal(actual_result, S_result), 'Arrays are not equal'
        print ("Building similarity matrix is successful!")
        print ("Context 2 Question attention")
        m_r, beta_r, q2c_r = sess.run([m, beta, q2c])
        output_r = sess.run(output)

[((201, 212), 'numpy.exp', 'np.exp', (['arr'], {}), '(arr)\n', (207, 212), True, 'import numpy as np\n'), ((1260, 1284), 'tensorflow.add', 'tf.add', (['logits', 'exp_mask'], {}), '(logits, exp_mask)\n', (1266, 1284), True, 'import tensorflow as tf\n'), ((1347, 1380), 'tensorflow.nn.softmax', 'tf.nn.softmax', (['masked_logits', 'dim'], {}), '(masked_logits, dim)\n', (1360, 1380), True, 'import tensorflow as tf\n'), ((1481, 1524), 'tensorflow.get_variable', 'tf.get_variable', (['"""w_sim_1"""'], {'initializer': 'w_1'}), "('w_sim_1', initializer=w_1)\n", (1496, 1524), True, 'import tensorflow as tf\n'), ((1555, 1598), 'tensorflow.get_variable', 'tf.get_variable', (['"""w_sim_2"""'], {'initializer': 'w_2'}), "('w_sim_2', initializer=w_2)\n", (1570, 1598), True, 'import tensorflow as tf\n'), ((1644, 1687), 'tensorflow.get_variable', 'tf.get_variable', (['"""w_sim_3"""'], {'initializer': 'w_3'}), "('w_sim_3', initializer=w_3)\n", (1659, 1687), True, 'import tensorflow as tf\n'), ((1816, 1850), 'tensorflow.transpose', 'tf.transpose', (['q_tile', '(1, 0, 3, 2)'], {}), '(q_tile, (1, 0, 3, 2))\n', (1828, 1850), True, 'import tensorflow as tf\n'), ((1884, 1912), 'tensorflow.expand_dims', 'tf.expand_dims', (['contexts', '(-1)'], {}), '(contexts, -1)\n', (1898, 1912), True, 'import tensorflow as tf\n'), ((2052, 2086), 'tensorflow.transpose', 'tf.transpose', (['result', '(0, 1, 3, 2)'], {}), '(result, (0, 1, 3, 2))\n', (2064, 2086), True, 'import tensorflow as tf\n'), ((2118, 2156), 'tensorflow.reshape', 'tf.reshape', (['result', '(-1, N * M, 2 * H)'], {}), '(result, (-1, N * M, 2 * H))\n', (2128, 2156), True, 'import tensorflow as tf\n'), ((2522, 2548), 'tensorflow.reshape', 'tf.reshape', (['term1', '(-1, N)'], {}), '(term1, (-1, N))\n', (2532, 2548), True, 'import tensorflow as tf\n'), ((2659, 2685), 'tensorflow.reshape', 'tf.reshape', (['term2', '(-1, M)'], {}), '(term2, (-1, M))\n', (2669, 2685), True, 'import tensorflow as tf\n'), ((2789, 2818), 'tensorflow.reshape', 'tf.reshape', (['term3', '(-1, N, M)'], {}), '(term3, (-1, N, M))\n', (2799, 2818), True, 'import tensorflow as tf\n'), ((2969, 2991), 'numpy.array', 'np.array', (['[True, True]'], {}), '([True, True])\n', (2977, 2991), True, 'import numpy as np\n'), ((3021, 3050), 'numpy.array', 'np.array', (['[True, True, False]'], {}), '([True, True, False])\n', (3029, 3050), True, 'import numpy as np\n'), ((3079, 3109), 'numpy.tile', 'np.tile', (['context_mask', '[BS, 1]'], {}), '(context_mask, [BS, 1])\n', (3086, 3109), True, 'import numpy as np\n'), ((3130, 3161), 'numpy.tile', 'np.tile', (['question_mask', '[BS, 1]'], {}), '(question_mask, [BS, 1])\n', (3137, 3161), True, 'import numpy as np\n'), ((3181, 3238), 'tensorflow.get_variable', 'tf.get_variable', (['"""context_mask"""'], {'initializer': 'context_mask'}), "('context_mask', initializer=context_mask)\n", (3196, 3238), True, 'import tensorflow as tf\n'), ((3259, 3318), 'tensorflow.get_variable', 'tf.get_variable', (['"""question_mask"""'], {'initializer': 'question_mask'}), "('question_mask', initializer=question_mask)\n", (3274, 3318), True, 'import tensorflow as tf\n'), ((3338, 3370), 'tensorflow.expand_dims', 'tf.expand_dims', (['context_mask', '(-1)'], {}), '(context_mask, -1)\n', (3352, 3370), True, 'import tensorflow as tf\n'), ((3404, 3437), 'tensorflow.expand_dims', 'tf.expand_dims', (['question_mask', '(-1)'], {}), '(question_mask, -1)\n', (3418, 3437), True, 'import tensorflow as tf\n'), ((3471, 3509), 'tensorflow.transpose', 'tf.transpose', (['question_mask', '(0, 2, 1)'], {}), '(question_mask, (0, 2, 1))\n', (3483, 3509), True, 'import tensorflow as tf\n'), ((3781, 3808), 'tensorflow.matmul', 'tf.matmul', (['alpha', 'questions'], {}), '(alpha, questions)\n', (3790, 3808), True, 'import tensorflow as tf\n'), ((4051, 4080), 'tensorflow.transpose', 'tf.transpose', (['beta', '(0, 2, 1)'], {}), '(beta, (0, 2, 1))\n', (4063, 4080), True, 'import tensorflow as tf\n'), ((4091, 4116), 'tensorflow.matmul', 'tf.matmul', (['beta', 'contexts'], {}), '(beta, contexts)\n', (4100, 4116), True, 'import tensorflow as tf\n'), ((4186, 4209), 'tensorflow.tile', 'tf.tile', (['q2c', '(1, N, 1)'], {}), '(q2c, (1, N, 1))\n', (4193, 4209), True, 'import tensorflow as tf\n'), ((4223, 4252), 'tensorflow.concat', 'tf.concat', (['[c2q, q2c]'], {'axis': '(2)'}), '([c2q, q2c], axis=2)\n', (4232, 4252), True, 'import tensorflow as tf\n'), ((4360, 4390), 'numpy.array', 'np.array', (['[1.0, 2.0, 3.0, 4.0]'], {}), '([1.0, 2.0, 3.0, 4.0])\n', (4368, 4390), True, 'import numpy as np\n'), ((4397, 4427), 'numpy.array', 'np.array', (['[5.0, 6.0, 7.0, 8.0]'], {}), '([5.0, 6.0, 7.0, 8.0])\n', (4405, 4427), True, 'import numpy as np\n'), ((4434, 4468), 'numpy.array', 'np.array', (['[13.0, 12.0, 11.0, 10.0]'], {}), '([13.0, 12.0, 11.0, 10.0])\n', (4442, 4468), True, 'import numpy as np\n'), ((4474, 4530), 'numpy.array', 'np.array', (['[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]'], {}), '([[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]])\n', (4482, 4530), True, 'import numpy as np\n'), ((4545, 4630), 'numpy.array', 'np.array', (['[[[1.0, 2.0, 3.0, 0.0], [5.0, 6.0, 7.0, 4.0], [8.0, 9.0, 10.0, 11.0]]]'], {}), '([[[1.0, 2.0, 3.0, 0.0], [5.0, 6.0, 7.0, 4.0], [8.0, 9.0, 10.0, 11.0]]]\n    )\n', (4553, 4630), True, 'import numpy as np\n'), ((4637, 4659), 'numpy.tile', 'np.tile', (['c', '[BS, 1, 1]'], {}), '(c, [BS, 1, 1])\n', (4644, 4659), True, 'import numpy as np\n'), ((4668, 4690), 'numpy.tile', 'np.tile', (['q', '[BS, 1, 1]'], {}), '(q, [BS, 1, 1])\n', (4675, 4690), True, 'import numpy as np\n'), ((4709, 4752), 'tensorflow.get_variable', 'tf.get_variable', (['"""questions"""'], {'initializer': 'q'}), "('questions', initializer=q)\n", (4724, 4752), True, 'import tensorflow as tf\n'), ((4768, 4810), 'tensorflow.get_variable', 'tf.get_variable', (['"""contexts"""'], {'initializer': 'c'}), "('contexts', initializer=c)\n", (4783, 4810), True, 'import tensorflow as tf\n'), ((5048, 5081), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), '()\n', (5079, 5081), True, 'import tensorflow as tf\n'), ((121, 140), 'numpy.max', 'np.max', (['arr'], {'axis': '(2)'}), '(arr, axis=2)\n', (127, 140), True, 'import numpy as np\n'), ((262, 287), 'numpy.sum', 'np.sum', (['exp_array'], {'axis': '(2)'}), '(exp_array, axis=2)\n', (268, 287), True, 'import numpy as np\n'), ((1740, 1768), 'tensorflow.expand_dims', 'tf.expand_dims', (['questions', '(0)'], {}), '(questions, 0)\n', (1754, 1768), True, 'import tensorflow as tf\n'), ((2002, 2018), 'tensorflow.shape', 'tf.shape', (['result'], {}), '(result)\n', (2010, 2018), True, 'import tensorflow as tf\n'), ((2195, 2211), 'tensorflow.shape', 'tf.shape', (['result'], {}), '(result)\n', (2203, 2211), True, 'import tensorflow as tf\n'), ((2435, 2472), 'tensorflow.reshape', 'tf.reshape', (['contexts', '(BS * N, 2 * H)'], {}), '(contexts, (BS * N, 2 * H))\n', (2445, 2472), True, 'import tensorflow as tf\n'), ((2472, 2499), 'tensorflow.expand_dims', 'tf.expand_dims', (['w_sim_1', '(-1)'], {}), '(w_sim_1, -1)\n', (2486, 2499), True, 'import tensorflow as tf\n'), ((2571, 2609), 'tensorflow.reshape', 'tf.reshape', (['questions', '(BS * M, 2 * H)'], {}), '(questions, (BS * M, 2 * H))\n', (2581, 2609), True, 'import tensorflow as tf\n'), ((2609, 2636), 'tensorflow.expand_dims', 'tf.expand_dims', (['w_sim_2', '(-1)'], {}), '(w_sim_2, -1)\n', (2623, 2636), True, 'import tensorflow as tf\n'), ((2708, 2747), 'tensorflow.reshape', 'tf.reshape', (['result', '(BS * N * M, 2 * H)'], {}), '(result, (BS * N * M, 2 * H))\n', (2718, 2747), True, 'import tensorflow as tf\n'), ((2748, 2775), 'tensorflow.expand_dims', 'tf.expand_dims', (['w_sim_3', '(-1)'], {}), '(w_sim_3, -1)\n', (2762, 2775), True, 'import tensorflow as tf\n'), ((2879, 2908), 'tensorflow.reshape', 'tf.reshape', (['term2', '(-1, 1, M)'], {}), '(term2, (-1, 1, M))\n', (2889, 2908), True, 'import tensorflow as tf\n'), ((3548, 3585), 'tensorflow.cast', 'tf.cast', (['context_mask'], {'dtype': 'tf.int32'}), '(context_mask, dtype=tf.int32)\n', (3555, 3585), True, 'import tensorflow as tf\n'), ((3599, 3637), 'tensorflow.cast', 'tf.cast', (['question_mask'], {'dtype': 'tf.int32'}), '(question_mask, dtype=tf.int32)\n', (3606, 3637), True, 'import tensorflow as tf\n'), ((4005, 4021), 'tensorflow.nn.softmax', 'tf.nn.softmax', (['m'], {}), '(m)\n', (4018, 4021), True, 'import tensorflow as tf\n'), ((4273, 4289), 'tensorflow.shape', 'tf.shape', (['output'], {}), '(output)\n', (4281, 4289), True, 'import tensorflow as tf\n'), ((5091, 5103), 'tensorflow.Session', 'tf.Session', ([], {}), '()\n', (5101, 5103), True, 'import tensorflow as tf\n'), ((5308, 5347), 'numpy.array_equal', 'np.array_equal', (['actual_result', 'S_result'], {}), '(actual_result, S_result)\n', (5322, 5347), True, 'import numpy as np\n'), ((1139, 1163), 'tensorflow.cast', 'tf.cast', (['mask', '"""float64"""'], {}), "(mask, 'float64')\n", (1146, 1163), True, 'import tensorflow as tf\n'), ((2840, 2869), 'tensorflow.reshape', 'tf.reshape', (['term1', '(-1, N, 1)'], {}), '(term1, (-1, N, 1))\n', (2850, 2869), True, 'import tensorflow as tf\n'), ((3927, 3960), 'tensorflow.cast', 'tf.cast', (['S_mask'], {'dtype': 'tf.float64'}), '(S_mask, dtype=tf.float64)\n', (3934, 3960), True, 'import tensorflow as tf\n'), ((5232, 5279), 'numpy.array', 'np.array', (['[[228, 772, 1372], [548, 1828, 3140]]'], {}), '([[228, 772, 1372], [548, 1828, 3140]])\n', (5240, 5279), True, 'import numpy as np\n')]