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, 23, 9, 50), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', ({(9, 42, 9, 49): '7 * 7 * 256'}, {}), '(7 * 7 * 256)', False, 'from tensorflow import keras\n'), ((10, 25, 10, 58), 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', ({(10, 46, 10, 57): '(7, 7, 256)'}, {}), '((7, 7, 256))', False, 'from tensorflow import keras\n'), ((12, 24, 12, 55), 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', ({(12, 47, 12, 49): '10', (12, 51, 12, 54): '100'}, {}), '(10, 100)', False, 'from tensorflow import keras\n'), ((13, 23, 13, 50), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', ({(13, 42, 13, 49): '7 * 7 * 256'}, {}), '(7 * 7 * 256)', False, 'from tensorflow import keras\n'), ((15, 23, 16, 67), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (), '', False, 'from tensorflow import keras\n'), ((17, 26, 17, 59), 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((18, 28, 18, 52), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((20, 23, 21, 66), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (), '', False, 'from tensorflow import keras\n'), ((22, 26, 22, 59), 'tensorflow.keras.layers.BatchNormalization', 'keras.layers.BatchNormalization', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((23, 28, 23, 52), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((25, 25, 26, 65), 'tensorflow.keras.layers.Conv2DTranspose', 'keras.layers.Conv2DTranspose', (), '', False, 'from tensorflow import keras\n'), ((34, 18, 34, 40), 'tensorflow.squeeze', 'tf.squeeze', ({(34, 29, 34, 36): 'label_x', (34, 38, 34, 39): '1'}, {}), '(label_x, 1)', True, 'import tensorflow as tf\n'), ((38, 12, 38, 45), 'tensorflow.math.multiply', 'tf.math.multiply', ({(38, 29, 38, 35): 'feat_x', (38, 37, 38, 44): 'label_x'}, {}), '(feat_x, label_x)', True, 'import tensorflow as tf\n'), ((56, 24, 56, 55), 'tensorflow.keras.layers.Embedding', 'keras.layers.Embedding', ({(56, 47, 56, 49): '10', (56, 51, 56, 54): '100'}, {}), '(10, 100)', False, 'from tensorflow import keras\n'), ((57, 28, 57, 55), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', ({(57, 47, 57, 54): '28 * 28 * 1'}, {}), '(28 * 28 * 1)', False, 'from tensorflow import keras\n'), ((58, 29, 58, 62), 'tensorflow.keras.layers.Reshape', 'keras.layers.Reshape', ({(58, 50, 58, 61): '(28, 28, 1)'}, {}), '((28, 28, 1))', False, 'from tensorflow import keras\n'), ((59, 22, 60, 75), 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (), '', False, 'from tensorflow import keras\n'), ((61, 22, 61, 46), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((62, 22, 62, 47), 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', ({(62, 43, 62, 46): '0.3'}, {}), '(0.3)', False, 'from tensorflow import keras\n'), ((63, 22, 64, 51), 'tensorflow.keras.layers.Conv2D', 'keras.layers.Conv2D', (), '', False, 'from tensorflow import keras\n'), ((65, 22, 65, 46), 'tensorflow.keras.layers.LeakyReLU', 'keras.layers.LeakyReLU', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((66, 22, 66, 47), 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', ({(66, 43, 66, 46): '0.3'}, {}), '(0.3)', False, 'from tensorflow import keras\n'), ((67, 23, 67, 45), 'tensorflow.keras.layers.Flatten', 'keras.layers.Flatten', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((68, 19, 68, 40), 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', ({(68, 38, 68, 39): '1'}, {}), '(1)', False, 'from tensorflow import keras\n'), ((76, 12, 76, 48), 'tensorflow.math.multiply', 'tf.math.multiply', ({(76, 29, 76, 37): 'images_x', (76, 39, 76, 47): 'labels_x'}, {}), '(images_x, labels_x)', True, 'import tensorflow as tf\n')]