import tensorflow as tf
from tensorflow import keras


class CondGeneratorModel(keras.Model):
    def __init__(self):
        super(CondGeneratorModel, self).__init__()
        # Expand 7*7*128 features into a (7,7,128) tensor
        self.dense_1 = keras.layers.Dense(7*7*256)
        self.reshape_1 = keras.layers.Reshape((7, 7, 256))
        # Expand (10,) to (7,7,1)
        self.embedder = keras.layers.Embedding(10, 100)
        self.dense_2 = keras.layers.Dense(7*7*256)
        # From (7,7,256) to (7,7,128)
        self.convt_1 = keras.layers.Conv2DTranspose(
            128, (5, 5), strides=1, padding='same', use_bias=False)
        self.convt_bn_1 = keras.layers.BatchNormalization()
        self.convt_relu_1 = keras.layers.LeakyReLU()
        # From (7,7,128) to (14,14,64)
        self.convt_2 = keras.layers.Conv2DTranspose(
            64, (5, 5), strides=2, padding='same', use_bias=False)
        self.convt_bn_2 = keras.layers.BatchNormalization()
        self.convt_relu_2 = keras.layers.LeakyReLU()
        # From (14,14,64) to (28,28,1)
        self.convt_out = keras.layers.Conv2DTranspose(
            1, (5, 5), strides=2, padding='same', use_bias=False)

    def call(self, inputs):
        feat_x = inputs[0]
        label = inputs[2]
        # Expand label input to be the same as latent feature
        label_x = self.embedder(label)
        label_x = self.dense_2(label_x)
        label_x = tf.squeeze(label_x, 1)
        # Expand features to image channels
        feat_x = self.dense_1(feat_x)
        # Combine latent feature and label input
        x = tf.math.multiply(feat_x, label_x)
        x = self.reshape_1(x)
        # From (7,7,256) to (7,7,128)
        x = self.convt_1(x)
        x = self.convt_bn_1(x)
        x = self.convt_relu_1(x)
        # From (7,7,128) to (14,14,64)
        x = self.convt_2(x)
        x = self.convt_bn_2(x)
        x = self.convt_relu_2(x)
        # From (14,14,64) to (28,28,1)
        x = self.convt_out(x)
        return [x, None, label]


class CondDiscriminatorModel(keras.Model):
    def __init__(self):
        super(CondDiscriminatorModel, self).__init__()
        self.embedder = keras.layers.Embedding(10, 100)
        self.expand_layer = keras.layers.Dense(28*28*1)
        self.reshape_layer = keras.layers.Reshape((28, 28, 1))
        self.conv_1 = keras.layers.Conv2D(
            64, (5, 5), strides=2, padding='same', input_shape=(28, 28, 1))
        self.relu_1 = keras.layers.LeakyReLU()
        self.drop_1 = keras.layers.Dropout(0.3)
        self.conv_2 = keras.layers.Conv2D(
            128, (5, 5), strides=2, padding='same')
        self.relu_2 = keras.layers.LeakyReLU()
        self.drop_2 = keras.layers.Dropout(0.3)
        self.flatten = keras.layers.Flatten()
        self.out = keras.layers.Dense(1)

    def call(self, inputs):
        images_x = inputs[0]
        labels = inputs[2]
        labels_x = self.embedder(labels)
        labels_x = self.expand_layer(labels_x)
        labels_x = self.reshape_layer(labels_x)
        x = tf.math.multiply(images_x, labels_x)
        x = self.conv_1(x)
        x = self.relu_1(x)
        x = self.drop_1(x)
        x = self.conv_2(x)
        x = self.relu_2(x)
        x = self.drop_2(x)
        x = self.flatten(x)
        x = self.out(x)
        return x

[(9, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(7 * 7 * 256)'], {}), False, 'from tensorflow import keras\n'), (10, 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', (['(7, 7, 256)'], {}), False, 'from tensorflow import keras\n'), (12, 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', (['(10)', '(100)'], {}), False, 'from tensorflow import keras\n'), (13, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(7 * 7 * 256)'], {}), False, 'from tensorflow import keras\n'), (15, 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(128)', '(5, 5)'], {'strides': '(1)', 'padding': '"""same"""', 'use_bias': '(False)'}), False, 'from tensorflow import keras\n'), (17, 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ([], {}), False, 'from tensorflow import keras\n'), (18, 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), False, 'from tensorflow import keras\n'), (20, 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(64)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'use_bias': '(False)'}), False, 'from tensorflow import keras\n'), (22, 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ([], {}), False, 'from tensorflow import keras\n'), (23, 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), False, 'from tensorflow import keras\n'), (25, 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(1)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'use_bias': '(False)'}), False, 'from tensorflow import keras\n'), (34, 'tensorflow.squeeze', 'tf.squeeze', (['label_x', '(1)'], {}), True, 'import tensorflow as tf\n'), (38, 'tensorflow.math.multiply', 'tf.math.multiply', (['feat_x', 'label_x'], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', (['(10)', '(100)'], {}), False, 'from tensorflow import keras\n'), (57, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(28 * 28 * 1)'], {}), False, 'from tensorflow import keras\n'), (58, 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', (['(28, 28, 1)'], {}), False, 'from tensorflow import keras\n'), (59, 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (['(64)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'input_shape': '(28, 28, 1)'}), False, 'from tensorflow import keras\n'), (61, 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), False, 'from tensorflow import keras\n'), (62, 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', (['(0.3)'], {}), False, 'from tensorflow import keras\n'), (63, 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (['(128)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""'}), False, 'from tensorflow import keras\n'), (65, 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), False, 'from tensorflow import keras\n'), (66, 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', (['(0.3)'], {}), False, 'from tensorflow import keras\n'), (67, 'tensorflow.keras.layers.Flatten', 'keras.layers.Flatten', ([], {}), False, 'from tensorflow import keras\n'), (68, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(1)'], {}), False, 'from tensorflow import keras\n'), (76, 'tensorflow.math.multiply', 'tf.math.multiply', (['images_x', 'labels_x'], {}), True, 'import tensorflow as tf\n')]