import tensorflow as tf
from .utils import as_KerasModel

layers = tf.keras.layers


@as_KerasModel
class GAN(tf.keras.Model):
    def __init__(self, discriminator=None, generator=None, latent_dim=128):
        super(GAN, self).__init__()

        if discriminator is None:
            discriminator = self.default_discriminator()

        if generator is None:
            generator = self.default_generator()

        self.discriminator = discriminator
        self.generator = generator
        self.latent_dim = latent_dim

    def compile(self, d_optimizer, g_optimizer, loss_fn):
        super(GAN, self).compile()
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        self.loss_fn = loss_fn
        self.d_loss_metric = tf.keras.metrics.Mean(name="d_loss")
        self.g_loss_metric = tf.keras.metrics.Mean(name="g_loss")

    @property
    def metrics(self):
        return [self.d_loss_metric, self.g_loss_metric]

    def train_step(self, real_images):
        # Sample random points in the latent space
        batch_size = tf.shape(real_images)[0]
        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))

        # Decode them to fake images
        generated_images = self.generator(random_latent_vectors)

        # Combine them with real images
        combined_images = tf.concat([generated_images, real_images], axis=0)

        # Assemble labels discriminating real from fake images
        labels = tf.concat(
            [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
        )
        # Add random noise to the labels - important trick!
        labels += 0.05 * tf.random.uniform(tf.shape(labels))

        # Train the discriminator
        with tf.GradientTape() as tape:
            predictions = self.discriminator(combined_images)
            d_loss = self.loss_fn(labels, predictions)
        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply_gradients(
            zip(grads, self.discriminator.trainable_weights)
        )

        # Sample random points in the latent space
        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))

        # Assemble labels that say "all real images"
        misleading_labels = tf.zeros((batch_size, 1))

        # Train the generator (note that we should *not* update the weights
        # of the discriminator)!
        with tf.GradientTape() as tape:
            predictions = self.discriminator(self.generator(random_latent_vectors))
            g_loss = self.loss_fn(misleading_labels, predictions)
        grads = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))

        # Update metrics
        self.d_loss_metric.update_state(d_loss)
        self.g_loss_metric.update_state(g_loss)
        return {
            "d_loss": self.d_loss_metric.result(),
            "g_loss": self.g_loss_metric.result(),
        }

    def default_generator(self, latent_dim=128):
        return tf.keras.Sequential(
            [
                tf.keras.Input(shape=(latent_dim,)),
                layers.Dense(8 * 8 * 128),
                layers.Reshape((8, 8, 128)),
                layers.Conv2DTranspose(128, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Conv2DTranspose(256, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Conv2DTranspose(512, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Conv2D(3, kernel_size=5, padding="same", activation="sigmoid"),
            ],
            name="generator",
        )

    def default_discriminator(self):
        return tf.keras.Sequential(
            [
                tf.keras.Input(shape=(64, 64, 3)),
                layers.Conv2D(64, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
                layers.LeakyReLU(alpha=0.2),
                layers.Flatten(),
                layers.Dropout(0.2),
                layers.Dense(1, activation="sigmoid"),
            ],
            name="discriminator",
        )

tensorflow.concat tensorflow.keras.Input tensorflow.shape tensorflow.zeros tensorflow.ones tensorflow.random.normal tensorflow.keras.metrics.Mean tensorflow.GradientTape
[(27, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {'name': '"""d_loss"""'}), True, 'import tensorflow as tf\n'), (28, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {'name': '"""g_loss"""'}), True, 'import tensorflow as tf\n'), (37, 'tensorflow.random.normal', 'tf.random.normal', ([], {'shape': '(batch_size, self.latent_dim)'}), True, 'import tensorflow as tf\n'), (43, 'tensorflow.concat', 'tf.concat', (['[generated_images, real_images]'], {'axis': '(0)'}), True, 'import tensorflow as tf\n'), (62, 'tensorflow.random.normal', 'tf.random.normal', ([], {'shape': '(batch_size, self.latent_dim)'}), True, 'import tensorflow as tf\n'), (65, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, 1)'], {}), True, 'import tensorflow as tf\n'), (36, 'tensorflow.shape', 'tf.shape', (['real_images'], {}), True, 'import tensorflow as tf\n'), (53, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (69, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (47, 'tensorflow.ones', 'tf.ones', (['(batch_size, 1)'], {}), True, 'import tensorflow as tf\n'), (47, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, 1)'], {}), True, 'import tensorflow as tf\n'), (50, 'tensorflow.shape', 'tf.shape', (['labels'], {}), True, 'import tensorflow as tf\n'), (86, 'tensorflow.keras.Input', 'tf.keras.Input', ([], {'shape': '(latent_dim,)'}), True, 'import tensorflow as tf\n'), (103, 'tensorflow.keras.Input', 'tf.keras.Input', ([], {'shape': '(64, 64, 3)'}), True, 'import tensorflow as tf\n')]