import os,sys
import numpy as np
import unittest

from deepmd.env import tf

from deepmd.utils.network import embedding_net
from deepmd.common import get_activation_func

class TestGelu(tf.test.TestCase):
    def setUp (self) :
        self.places = 6
        self.sess = self.test_session().__enter__()
        self.inputs = tf.reshape(tf.constant([ 0., 1., 2., 3.], dtype = tf.float64), [-1, 1])
        self.refout = [[ 0.37703893, -0.38242253, -0.1862878,  -0.23220415,  2.28706995, -0.40754364,
                         0.22086098, -0.2690335 ],
                       [ 2.167494,    0.72560347,  0.99234317,  0.50832127,  5.20665818,  0.58361587,
                         1.57217107,  0.67395218],
                       [ 4.19655852,  2.04779208,  2.20239826,  1.69247695,  8.38305924,  1.69006845,
                         2.97176052,  1.76098426],
                       [ 6.21460216,  3.52613278,  3.39508271,  2.817003,   11.521799,    2.91028145,
                         4.41870371,  2.82610791]]       
        
    def test_activation_function_gelu_custom(self):
        network_size = [2, 4, 8]
        out = embedding_net(self.inputs, 
                            network_size, 
                            tf.float64,
                            activation_fn = get_activation_func('gelu'),
                            name_suffix = 'gelu_custom',
                            seed = 1, 
                            uniform_seed = True)
        self.sess.run(tf.global_variables_initializer())
        myout = self.sess.run(out)
        np.testing.assert_almost_equal(self.refout, myout, self.places)


    def test_activation_function_gelu_tensorflow(self):
        network_size = [2, 4, 8]
        out = embedding_net(self.inputs, 
                            network_size, 
                            tf.float64,
                            activation_fn = get_activation_func('gelu_tf'),
                            name_suffix = 'gelu_tensorflow',
                            seed = 1, 
                            uniform_seed = True)
        self.sess.run(tf.global_variables_initializer())
        myout = self.sess.run(out)
        np.testing.assert_almost_equal(self.refout, myout, self.places)


if __name__ == '__main__':
    unittest.main()
