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

[((250, 281), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(7 * 7 * 256)'], {}), '(7 * 7 * 256)\n', (268, 281), False, 'from tensorflow import keras\n'), ((303, 336), 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', (['(7, 7, 256)'], {}), '((7, 7, 256))\n', (323, 336), False, 'from tensorflow import keras\n'), ((395, 426), 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', (['(10)', '(100)'], {}), '(10, 100)\n', (417, 426), False, 'from tensorflow import keras\n'), ((450, 481), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(7 * 7 * 256)'], {}), '(7 * 7 * 256)\n', (468, 481), False, 'from tensorflow import keras\n'), ((539, 627), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(128)', '(5, 5)'], {'strides': '(1)', 'padding': '"""same"""', 'use_bias': '(False)'}), "(128, (5, 5), strides=1, padding='same',\n    use_bias=False)\n", (567, 627), False, 'from tensorflow import keras\n'), ((663, 696), 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ([], {}), '()\n', (694, 696), False, 'from tensorflow import keras\n'), ((725, 749), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), '()\n', (747, 749), False, 'from tensorflow import keras\n'), ((812, 899), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(64)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'use_bias': '(False)'}), "(64, (5, 5), strides=2, padding='same',\n    use_bias=False)\n", (840, 899), False, 'from tensorflow import keras\n'), ((935, 968), 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ([], {}), '()\n', (966, 968), False, 'from tensorflow import keras\n'), ((997, 1021), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), '()\n', (1019, 1021), False, 'from tensorflow import keras\n'), ((1086, 1173), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (['(1)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'use_bias': '(False)'}), "(1, (5, 5), strides=2, padding='same', use_bias\n    =False)\n", (1114, 1173), False, 'from tensorflow import keras\n'), ((1423, 1445), 'tensorflow.squeeze', 'tf.squeeze', (['label_x', '(1)'], {}), '(label_x, 1)\n', (1433, 1445), True, 'import tensorflow as tf\n'), ((1589, 1622), 'tensorflow.math.multiply', 'tf.math.multiply', (['feat_x', 'label_x'], {}), '(feat_x, label_x)\n', (1605, 1622), True, 'import tensorflow as tf\n'), ((2163, 2194), 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', (['(10)', '(100)'], {}), '(10, 100)\n', (2185, 2194), False, 'from tensorflow import keras\n'), ((2223, 2254), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(28 * 28 * 1)'], {}), '(28 * 28 * 1)\n', (2241, 2254), False, 'from tensorflow import keras\n'), ((2280, 2313), 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', (['(28, 28, 1)'], {}), '((28, 28, 1))\n', (2300, 2313), False, 'from tensorflow import keras\n'), ((2336, 2423), 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (['(64)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""', 'input_shape': '(28, 28, 1)'}), "(64, (5, 5), strides=2, padding='same', input_shape=(28,\n    28, 1))\n", (2355, 2423), False, 'from tensorflow import keras\n'), ((2455, 2479), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), '()\n', (2477, 2479), False, 'from tensorflow import keras\n'), ((2502, 2527), 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', (['(0.3)'], {}), '(0.3)\n', (2522, 2527), False, 'from tensorflow import keras\n'), ((2550, 2609), 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (['(128)', '(5, 5)'], {'strides': '(2)', 'padding': '"""same"""'}), "(128, (5, 5), strides=2, padding='same')\n", (2569, 2609), False, 'from tensorflow import keras\n'), ((2645, 2669), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ([], {}), '()\n', (2667, 2669), False, 'from tensorflow import keras\n'), ((2692, 2717), 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', (['(0.3)'], {}), '(0.3)\n', (2712, 2717), False, 'from tensorflow import keras\n'), ((2741, 2763), 'tensorflow.keras.layers.Flatten', 'keras.layers.Flatten', ([], {}), '()\n', (2761, 2763), False, 'from tensorflow import keras\n'), ((2783, 2804), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(1)'], {}), '(1)\n', (2801, 2804), False, 'from tensorflow import keras\n'), ((3038, 3074), 'tensorflow.math.multiply', 'tf.math.multiply', (['images_x', 'labels_x'], {}), '(images_x, labels_x)\n', (3054, 3074), True, 'import tensorflow as tf\n')]