from __future__ import absolute_import, print_function

import unittest

import os
import tensorflow as tf
from tensorflow.keras import regularizers

from niftynet.network.simple_gan import SimpleGAN
from tests.niftynet_testcase import NiftyNetTestCase

class SimpleGANTest(NiftyNetTestCase):
    def test_3d_reg_shape(self):
        input_shape = (2, 32, 32, 32, 1)
        noise_shape = (2, 512)
        x = tf.ones(input_shape)
        r = tf.ones(noise_shape)

        simple_gan_instance = SimpleGAN()
        out = simple_gan_instance(r, x, is_training=True)

        with self.cached_session() as sess:
            sess.run(tf.compat.v1.global_variables_initializer())
            out = sess.run(out)
            self.assertAllClose(input_shape, out[0].shape)
            self.assertAllClose((2, 1), out[1].shape)
            self.assertAllClose((2, 1), out[2].shape)

    def test_2d_reg_shape(self):
        input_shape = (2, 64, 64, 1)
        noise_shape = (2, 512)
        x = tf.ones(input_shape)
        r = tf.ones(noise_shape)

        simple_gan_instance = SimpleGAN()
        out = simple_gan_instance(r, x, is_training=True)

        with self.cached_session() as sess:
            sess.run(tf.compat.v1.global_variables_initializer())
            out = sess.run(out)
            self.assertAllClose(input_shape, out[0].shape)
            self.assertAllClose((2, 1), out[1].shape)
            self.assertAllClose((2, 1), out[2].shape)



if __name__ == "__main__":
    tf.test.main()

[((48, 4, 48, 18), 'tensorflow.test.main', 'tf.test.main', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((16, 12, 16, 32), 'tensorflow.ones', 'tf.ones', ({(16, 20, 16, 31): 'input_shape'}, {}), '(input_shape)', True, 'import tensorflow as tf\n'), ((17, 12, 17, 32), 'tensorflow.ones', 'tf.ones', ({(17, 20, 17, 31): 'noise_shape'}, {}), '(noise_shape)', True, 'import tensorflow as tf\n'), ((19, 30, 19, 41), 'niftynet.network.simple_gan.SimpleGAN', 'SimpleGAN', ({}, {}), '()', False, 'from niftynet.network.simple_gan import SimpleGAN\n'), ((32, 12, 32, 32), 'tensorflow.ones', 'tf.ones', ({(32, 20, 32, 31): 'input_shape'}, {}), '(input_shape)', True, 'import tensorflow as tf\n'), ((33, 12, 33, 32), 'tensorflow.ones', 'tf.ones', ({(33, 20, 33, 31): 'noise_shape'}, {}), '(noise_shape)', True, 'import tensorflow as tf\n'), ((35, 30, 35, 41), 'niftynet.network.simple_gan.SimpleGAN', 'SimpleGAN', ({}, {}), '()', False, 'from niftynet.network.simple_gan import SimpleGAN\n'), ((23, 21, 23, 64), 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((39, 21, 39, 64), 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n')]