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, 'tensorflow.test.main', 'tf.test.main', True, 'import tensorflow as tf\n'), (16, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (17, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (19, 'niftynet.network.simple_gan.SimpleGAN', 'SimpleGAN', False, 'from niftynet.network.simple_gan import SimpleGAN\n'), (32, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (33, 'tensorflow.ones', 'tf.ones', True, 'import tensorflow as tf\n'), (35, 'niftynet.network.simple_gan.SimpleGAN', 'SimpleGAN', False, 'from niftynet.network.simple_gan import SimpleGAN\n'), (23, 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', True, 'import tensorflow as tf\n'), (39, 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', True, 'import tensorflow as tf\n')]