import os
import sys
sys.path.append(os.getcwd())

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras import Model

import config

# Prepare the data
mnist = tf.keras.datasets.mnist

# Parameter setting for this experiment
IMAGE_SIZE = 784

DISCRIMINATOR_HIDDEN_SIZE = 128

GENERATOR_HIDDEN_SIZE = 256

Q_HIDDEN_SIZE = 128

CONDITIONAL_SIZE = 10

Z_DIM = 16

BATCH_SIZE = 32

# Data normalization and capsulation into a Dataset object
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train / 255.0

train_ds = iter(tf.data.Dataset.from_tensor_slices(x_train).repeat().batch(BATCH_SIZE))


# MLP-based Generator
class Generator(Model):
    def __init__(self):
        super(Generator, self).__init__()
        self.d1 = Dense(GENERATOR_HIDDEN_SIZE, activation="relu")
        self.d2 = Dense(IMAGE_SIZE, activation="sigmoid")

    def call(self, z: tf.Tensor, c: tf.Tensor) -> tf.Tensor:
        """Forward computation

        Arguments:
            z {tf.Tensor} -- the random noise for generation
            c {tf.Tensor} -- the conditional input feature, only support digit choice now

        Returns:
            tf.Tensor -- the generated sample
        """
        x = tf.concat([z, c], -1)
        x = self.d1(x)
        x = self.d2(x)
        return x


def generator_loss(fake_logit: tf.Tensor) -> tf.Tensor:
    """Computation of loss for generator training

    Arguments:
        fake_logit {tf.Tensor} -- the logit output of discriminator when feeded with generated samples

    Returns:
        tf.Tensor -- the loss of generator
    """
    return -tf.reduce_mean(tf.math.log(fake_logit + 1e-8))


# MLP-based Discriminator
class Discriminator(Model):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.d1 = Dense(DISCRIMINATOR_HIDDEN_SIZE, activation="relu")
        self.d2 = Dense(1, activation="sigmoid")

    def call(self, x: tf.Tensor) -> tf.Tensor:
        """Forward computation

        Arguments:
            x {tf.Tensor} -- the input of discriminator, should be real samples or generated samples

        Returns:
            tf.Tensor -- the binary classification result
        """
        x = self.d1(x)
        x = self.d2(x)
        return x


def discriminator_loss(real_logit: tf.Tensor, fake_logit: tf.Tensor) -> tf.Tensor:
    """Computation of loss for discrimination training

    Arguments:
        real_logit {tf.Tensor} -- the logit output of real samples
        fake_logit {tf.Tensor} -- the logit output of generated samples

    Returns:
        tf.Tensor -- the loss of discriminator
    """
    return -tf.reduce_mean(tf.math.log(real_logit + 1e-8) + tf.math.log(1 - fake_logit + 1e-8))


# MLP-based Q-value network, for mutual information induction
class Q(Model):
    def __init__(self):
        super(Q, self).__init__()
        self.d1 = Dense(Q_HIDDEN_SIZE, activation="relu")
        self.d2 = Dense(CONDITIONAL_SIZE, activation="softmax")

    def call(self, x: tf.Tensor) -> tf.Tensor:
        """Forward computation

        Arguments:
            x {tf.Tensor} -- the input samples, generated by generator

        Returns:
            tf.Tensor -- the approximated conditional information (for mutual information computation)
        """
        x = self.d1(x)
        x = self.d2(x)
        return x


def Q_loss(Q_c_given_x: tf.Tensor, c: tf.Tensor) -> tf.Tensor:
    """Computation of loss for Q-value network training

    Arguments:
        Q_c_given_x {tf.Tensor} -- the result of Q-value network, trying to decode the conditional information based on the generated samples
        c {tf.Tensor} -- the actual conditional information

    Returns:
        tf.Tensor -- the loss of Q-value network
    """
    return tf.reduce_mean(-tf.reduce_sum(tf.math.log(Q_c_given_x + 1e-8) * c, 1))


def plot(samples: np.ndarray) -> plt.figure:
    """Result observation from generator

    Arguments:
        samples {np.ndarray} -- the output of generator, shape => (batch_size, X_dim)

    Returns:
        plt.figure -- the figure created by matplotlib
    """
    fig = plt.figure(figsize=(4, 4))
    gs = gridspec.GridSpec(4, 4)
    gs.update(wspace=0.05, hspace=0.05)
    for i, sample in enumerate(samples):
        ax = plt.subplot(gs[i])
        plt.axis("off")
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_aspect("equal")
        plt.imshow(sample.reshape(28, 28), cmap="Greys_r")
    return fig


def sample_Z(m: int, n: int) -> np.ndarray:
    """Sampling algorithm of noise input

    Arguments:
        m {int} -- Batch size
        n {int} -- Noise input dimension

    Returns:
        np.ndarray -- a random-initialized numpy array with shape = (m ,n)
    """
    return np.random.uniform(-1.0, 1.0, size=[m, n]).astype(np.float32)


def sample_c(m: int) -> np.ndarray:
    """Randomly selected conditional information

    Arguments:
        m {int} -- Batch size

    Returns:
        np.ndarray -- the conditional information, for example, trying to generate '1' with [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    """
    return np.random.multinomial(1, 10 * [0.1], size=m).astype(np.float32)


# @tf.function
def discriminator_train_step(generator: Generator, discriminator: Discriminator, X: tf.Tensor, Z: tf.Tensor, c: tf.Tensor, 
                             optimizer: tf.keras.optimizers.Optimizer) -> tf.Tensor:
    """Single training step of discriminator

    Arguments:
        generator {Generator} -- the generator instance
        discriminator {Discriminator} -- the discriminator instance
        X {tf.Tensor} -- the real sample input
        Z {tf.Tensor} -- the random noise
        c {tf.Tensor} -- the conditional information
        optimizer {tf.keras.optimizers.Optimizer} -- the optimizer for bp update

    Returns:
        tf.Tensor -- the discriminator loss of current step
    """
    with tf.GradientTape() as tape:
        G_sample = generator(Z, c)
        d_real = discriminator(X)
        d_fake = discriminator(G_sample)
        loss = discriminator_loss(d_real, d_fake)
    gradients = tape.gradient(loss, discriminator.trainable_variables)
    optimizer.apply_gradients(zip(gradients, discriminator.trainable_variables))
    return loss


# @tf.function
def generator_train_step(generator: Generator, discriminator: Discriminator, Z: tf.Tensor, c: tf.Tensor, 
                         optimizer: tf.keras.optimizers.Optimizer) -> tf.Tensor:
    """Single training step of generator

    Arguments:
        generator {Generator} -- the generator instance
        discriminator {Discriminator} -- the discriminator instance
        Z {tf.Tensor} -- the random noise for sample generation
        c {tf.Tensor} -- the conditional information
        optimizer {tf.keras.optimizers.Optimizer} -- the optimizer for bp update

    Returns:
        tf.Tensor -- the generator loss of current step
    """
    with tf.GradientTape() as tape:
        G_sample = generator(Z, c)
        d_fake = discriminator(G_sample)
        loss = generator_loss(d_fake)
    gradients = tape.gradient(loss, generator.trainable_variables)
    optimizer.apply_gradients(zip(gradients, generator.trainable_variables))
    return loss


# @tf.function
def Q_train_step(generator: Generator, q: Q, Z: tf.Tensor, c: tf.Tensor, optimizer: tf.keras.optimizers.Optimizer) -> tf.Tensor:
    """Single training step of Q-value network training

    Arguments:
        generator {Generator} -- the generator instance
        q {Q} -- the Q-value network instance
        Z {tf.Tensor} -- the random noise of sample generation
        c {tf.Tensor} -- the conditional information
        optimizer {tf.keras.optimizers.Optimizer} -- the optimizer for bp update

    Returns:
        tf.Tensor -- the loss of Q-value network for current step
    """
    with tf.GradientTape() as tape:
        G_sample = generator(Z, c)
        Q_c_given_x = q(G_sample)
        loss = Q_loss(Q_c_given_x, c)
    gradients = tape.gradient(loss, q.trainable_variables)
    optimizer.apply_gradients(zip(gradients, q.trainable_variables))
    return loss


def main():
    # Instantiation of generator, discriminator, and Q-value network
    generator = Generator()
    discriminator = Discriminator()
    q = Q()

    # Set up three different optimizers for each network separately, maybe it is ok to share the same one
    generator_solver = tf.keras.optimizers.Adam()
    discriminator_solver = tf.keras.optimizers.Adam()
    q_solver = tf.keras.optimizers.Adam()

    out_path = os.path.join(config.OUTPUT_DIR, "infoGAN")
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    # Training starts!
    i = 0
    for it in range(1000000):
        if it % 1000 == 0:
            n_sample = 16
            Z_noise = sample_Z(n_sample, Z_DIM)
            idx = np.random.randint(0, CONDITIONAL_SIZE)
            c_noise = np.zeros([n_sample, CONDITIONAL_SIZE])
            c_noise[range(n_sample), idx] = 1
            Z_noise = tf.constant(Z_noise)
            c_noise = tf.constant(c_noise.astype(np.float32))
            samples = generator(Z_noise, c_noise)
            samples = samples.numpy()
            fig = plot(samples)
            file_path = os.path.join(out_path, "{}.png".format(str(i).zfill(3)))
            plt.savefig(file_path, bbox_inches="tight")
            i += 1
            plt.close(fig)

        # read data
        images = next(train_ds)
        images = tf.reshape(images, [BATCH_SIZE, -1])
        
        # get noise and conditional information
        Z_noise = tf.constant(sample_Z(BATCH_SIZE, Z_DIM))
        c_noise = tf.constant(sample_c(BATCH_SIZE))

        # train steps
        d_loss = discriminator_train_step(generator, discriminator, images, Z_noise, c_noise, discriminator_solver)
        g_loss = generator_train_step(generator, discriminator, Z_noise, c_noise, generator_solver)
        _ = Q_train_step(generator, q, Z_noise, c_noise, q_solver)

        if it % 1000 == 0:
            print("Iter: {}".format(it))
            print("D loss: {:.4}".format(float(d_loss)))
            print("G loss: {:.4}".format(float(g_loss)))
            print()


if __name__ == "__main__":
    main()
