#!/usr/bin/env python
# coding: utf-8

# In[1]:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import pickle
import time
import random


# In[8]:


import PIL
from PIL import Image
from tensorflow import keras
import keras.backend as K
import tensorflow as tf
from keras.optimizers import Adam
from keras.models import Sequential
from keras import layers,Model,Input
from keras.layers import Lambda,Reshape,UpSampling2D,ReLU,add,ZeroPadding2D
from keras.layers import Activation,BatchNormalization,Concatenate,concatenate
from keras.layers import Dense,Conv2D,Flatten,Dropout,LeakyReLU
from keras.preprocessing.image import ImageDataGenerator


# ### Conditioning Augmentation Network

# In[3]:


# conditioned by the text.
def conditioning_augmentation(x):
    """The mean_logsigma passed as argument is converted into the text conditioning variable.

    Args:
        x: The output of the text embedding passed through a FC layer with LeakyReLU non-linearity.

    Returns:
        c: The text conditioning variable after computation.
    """
    mean = x[:, :128]
    log_sigma = x[:, 128:]

    stddev = tf.math.exp(log_sigma)
    epsilon = K.random_normal(shape=K.constant((mean.shape[1], ), dtype='int32'))
    c = mean + stddev * epsilon
    return c

def build_ca_network():
    """Builds the conditioning augmentation network.
    """
    input_layer1 = Input(shape=(1024,)) #size of the vocabulary in the text data
    mls = Dense(256)(input_layer1)
    mls = LeakyReLU(alpha=0.2)(mls)
    ca = Lambda(conditioning_augmentation)(mls)
    return Model(inputs=[input_layer1], outputs=[ca]) 


# ### Stage 1 Generator Network

# In[4]:


def UpSamplingBlock(x, num_kernels):
    """An Upsample block with Upsampling2D, Conv2D, BatchNormalization and a ReLU activation.

    Args:
        x: The preceding layer as input.
        num_kernels: Number of kernels for the Conv2D layer.

    Returns:
        x: The final activation layer after the Upsampling block.
    """
    x = UpSampling2D(size=(2,2))(x)
    x = Conv2D(num_kernels, kernel_size=(3,3), padding='same', strides=1, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x) #prevent from mode collapse
    x = ReLU()(x)
    return x


def build_stage1_generator():

    input_layer1 = Input(shape=(1024,))
    ca = Dense(256)(input_layer1)
    ca = LeakyReLU(alpha=0.2)(ca)

    # Obtain the conditioned text
    c = Lambda(conditioning_augmentation)(ca)

    input_layer2 = Input(shape=(100,))
    concat = Concatenate(axis=1)([c, input_layer2]) 

    x = Dense(16384, use_bias=False)(concat) 
    x = ReLU()(x)
    x = Reshape((4, 4, 1024), input_shape=(16384,))(x)

    x = UpSamplingBlock(x, 512) 
    x = UpSamplingBlock(x, 256)
    x = UpSamplingBlock(x, 128)
    x = UpSamplingBlock(x, 64)   # upsampled our image to 64*64*3 

    x = Conv2D(3, kernel_size=3, padding='same', strides=1, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = Activation('tanh')(x)

    stage1_gen = Model(inputs=[input_layer1, input_layer2], outputs=[x, ca]) 
    return stage1_gen


# In[5]:


generator = build_stage1_generator()
generator.summary()


# ### Stage 1 Discriminator Network

# In[9]:


def ConvBlock(x, num_kernels, kernel_size=(4,4), strides=2, activation=True):
    """A ConvBlock with a Conv2D, BatchNormalization and LeakyReLU activation.

    Args:
        x: The preceding layer as input.
        num_kernels: Number of kernels for the Conv2D layer.

    Returns:
        x: The final activation layer after the ConvBlock block.
    """
    x = Conv2D(num_kernels, kernel_size=kernel_size, padding='same', strides=strides, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    
    if activation:
        x = LeakyReLU(alpha=0.2)(x)
    return x


def build_embedding_compressor():
    """Build embedding compressor model
    """
    input_layer1 = Input(shape=(1024,)) 
    x = Dense(128)(input_layer1)
    x = ReLU()(x)

    model = Model(inputs=[input_layer1], outputs=[x])
    return model

# the discriminator is fed with two inputs, the feature from Generator and the text embedding
def build_stage1_discriminator():
    """Builds the Stage 1 Discriminator that uses the 64x64 resolution images from the generator
    and the compressed and spatially replicated embedding.

    Returns:
        Stage 1 Discriminator Model for StackGAN.
    """
    input_layer1 = Input(shape=(64, 64, 3))  

    x = Conv2D(64, kernel_size=(4,4), strides=2, padding='same', use_bias=False,
                kernel_initializer='he_uniform')(input_layer1)
    x = LeakyReLU(alpha=0.2)(x)

    x = ConvBlock(x, 128)
    x = ConvBlock(x, 256)
    x = ConvBlock(x, 512)

    # Obtain the compressed and spatially replicated text embedding
    input_layer2 = Input(shape=(4, 4, 128)) #2nd input to discriminator, text embedding
    concat = concatenate([x, input_layer2])

    x1 = Conv2D(512, kernel_size=(1,1), padding='same', strides=1, use_bias=False,
                kernel_initializer='he_uniform')(concat)
    x1 = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    x1 = LeakyReLU(alpha=0.2)(x)

    # Flatten and add a FC layer to predict.
    x1 = Flatten()(x1)
    x1 = Dense(1)(x1)
    x1 = Activation('sigmoid')(x1)

    stage1_dis = Model(inputs=[input_layer1, input_layer2], outputs=[x1])  
    return stage1_dis


# In[10]:


discriminator = build_stage1_discriminator()
discriminator.summary()


# ### Stage 1 Adversarial Model (Building a GAN)

# In[11]:


# Building GAN with Generator and Discriminator

def build_adversarial(generator_model, discriminator_model):
    """Stage 1 Adversarial model.

    Args:
        generator_model: Stage 1 Generator Model
        discriminator_model: Stage 1 Discriminator Model

    Returns:
        Adversarial Model.
    """
    input_layer1 = Input(shape=(1024,))  
    input_layer2 = Input(shape=(100,)) 
    input_layer3 = Input(shape=(4, 4, 128)) 

    x, ca = generator_model([input_layer1, input_layer2]) #text,noise

    discriminator_model.trainable = False 

    probabilities = discriminator_model([x, input_layer3]) 
    adversarial_model = Model(inputs=[input_layer1, input_layer2, input_layer3], outputs=[probabilities, ca])
    return adversarial_model


# In[12]:


ganstage1 = build_adversarial(generator, discriminator)
ganstage1.summary()


# ### Train Utilities

# In[13]:


def checkpoint_prefix():
    checkpoint_dir = './training_checkpoints'
    checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')

    return checkpoint_prefix

def adversarial_loss(y_true, y_pred):
    mean = y_pred[:, :128]
    ls = y_pred[:, 128:]
    loss = -ls + 0.5 * (-1 + tf.math.exp(2.0 * ls) + tf.math.square(mean))
    loss = K.mean(loss)
    return loss

def normalize(input_image, real_image):
    input_image = (input_image / 127.5) - 1
    real_image = (real_image / 127.5) - 1

    return input_image, real_image

def load_class_ids_filenames(class_id_path, filename_path):
    with open(class_id_path, 'rb') as file:
        class_id = pickle.load(file, encoding='latin1')

    with open(filename_path, 'rb') as file:
        filename = pickle.load(file, encoding='latin1')

    return class_id, filename

def load_text_embeddings(text_embeddings):
    with open(text_embeddings, 'rb') as file:
        embeds = pickle.load(file, encoding='latin1')
        embeds = np.array(embeds)

    return embeds

def load_bbox(data_path):
    bbox_path = data_path + '/bounding_boxes.txt'
    image_path = data_path + '/images.txt'
    bbox_df = pd.read_csv(bbox_path, delim_whitespace=True, header=None).astype(int)
    filename_df = pd.read_csv(image_path, delim_whitespace=True, header=None)

    filenames = filename_df[1].tolist()
    bbox_dict = {i[:-4]:[] for i in filenames[:2]}

    for i in range(0, len(filenames)):
        bbox = bbox_df.iloc[i][1:].tolist()
        dict_key = filenames[i][:-4]
        bbox_dict[dict_key] = bbox

    return bbox_dict

def load_images(image_path, bounding_box, size):
    """Crops the image to the bounding box and then resizes it.
    """
    image = Image.open(image_path).convert('RGB')
    w, h = image.size
    if bounding_box is not None:
        r = int(np.maximum(bounding_box[2], bounding_box[3]) * 0.75)
        c_x = int((bounding_box[0] + bounding_box[2]) / 2)
        c_y = int((bounding_box[1] + bounding_box[3]) / 2)
        y1 = np.maximum(0, c_y - r)
        y2 = np.minimum(h, c_y + r)
        x1 = np.maximum(0, c_x - r)
        x2 = np.minimum(w, c_x + r)
        image = image.crop([x1, y1, x2, y2])

    image = image.resize(size, PIL.Image.BILINEAR)
    return image

def load_data(filename_path, class_id_path, dataset_path, embeddings_path, size):
    """Loads the Dataset.
    """
    data_dir = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/birds"
    train_dir = data_dir + "/train"
    test_dir = data_dir + "/test"
    embeddings_path_train = train_dir + "/char-CNN-RNN-embeddings.pickle"
    embeddings_path_test = test_dir + "/char-CNN-RNN-embeddings.pickle"
    filename_path_train = train_dir + "/filenames.pickle"
    filename_path_test = test_dir + "/filenames.pickle"
    class_id_path_train = train_dir + "/class_info.pickle"
    class_id_path_test = test_dir + "/class_info.pickle"
    dataset_path = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/CUB_200_2011"
    class_id, filenames = load_class_ids_filenames(class_id_path, filename_path)
    embeddings = load_text_embeddings(embeddings_path)
    bbox_dict = load_bbox(dataset_path)

    x, y, embeds = [], [], []

    for i, filename in enumerate(filenames):
        bbox = bbox_dict[filename]

        try:
            image_path = f'{dataset_path}/images/{filename}.jpg'
            image = load_images(image_path, bbox, size)
            e = embeddings[i, :, :]
            embed_index = np.random.randint(0, e.shape[0] - 1)
            embed = e[embed_index, :]

            x.append(np.array(image))
            y.append(class_id[i])
            embeds.append(embed)

        except Exception as e:
            print(f'{e}')
    
    x = np.array(x)
    y = np.array(y)
    embeds = np.array(embeds)
    
    return x, y, embeds

def save_image(file, save_path):
    """Saves the image at the specified file path.
    """
    image = plt.figure()
    ax = image.add_subplot(1,1,1)
    ax.imshow(file)
    ax.axis("off")
    plt.savefig(save_path)


# In[28]:


############################################################
# StackGAN class
############################################################

class StackGanStage1(object):
    """StackGAN Stage 1 class."""

    data_dir = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/birds"
    train_dir = data_dir + "/train"
    test_dir = data_dir + "/test"
    embeddings_path_train = train_dir + "/char-CNN-RNN-embeddings.pickle"
    embeddings_path_test = test_dir + "/char-CNN-RNN-embeddings.pickle"
    filename_path_train = train_dir + "/filenames.pickle"
    filename_path_test = test_dir + "/filenames.pickle"
    class_id_path_train = train_dir + "/class_info.pickle"
    class_id_path_test = test_dir + "/class_info.pickle"
    dataset_path = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/CUB_200_2011"
    def __init__(self, epochs=500, z_dim=100, batch_size=64, enable_function=True, stage1_generator_lr=0.0002, stage1_discriminator_lr=0.0002):
        self.epochs = epochs
        self.z_dim = z_dim
        self.enable_function = enable_function
        self.stage1_generator_lr = stage1_generator_lr
        self.stage1_discriminator_lr = stage1_discriminator_lr
        self.image_size = 64
        self.conditioning_dim = 128
        self.batch_size = batch_size
        
        self.stage1_generator_optimizer = Adam(lr=stage1_generator_lr, beta_1=0.5, beta_2=0.999)
        self.stage1_discriminator_optimizer = Adam(lr=stage1_discriminator_lr, beta_1=0.5, beta_2=0.999)
        
        self.stage1_generator = build_stage1_generator()
        self.stage1_generator.compile(loss='mse', optimizer=self.stage1_generator_optimizer)

        self.stage1_discriminator = build_stage1_discriminator()
        self.stage1_discriminator.compile(loss='binary_crossentropy', optimizer=self.stage1_discriminator_optimizer)

        self.ca_network = build_ca_network()
        self.ca_network.compile(loss='binary_crossentropy', optimizer='Adam')

        self.embedding_compressor = build_embedding_compressor()
        self.embedding_compressor.compile(loss='binary_crossentropy', optimizer='Adam')

        self.stage1_adversarial = build_adversarial(self.stage1_generator, self.stage1_discriminator)
        self.stage1_adversarial.compile(loss=['binary_crossentropy', adversarial_loss], loss_weights=[1, 2.0], optimizer=self.stage1_generator_optimizer)

        self.checkpoint1 = tf.train.Checkpoint(
            generator_optimizer=self.stage1_generator_optimizer,
            discriminator_optimizer=self.stage1_discriminator_optimizer,
            generator=self.stage1_generator,
            discriminator=self.stage1_discriminator)

    def visualize_stage1(self):
        """Running Tensorboard visualizations.
        """
        tb = TensorBoard(log_dir="logs/".format(time.time()))
        tb.set_model(self.stage1_generator)
        tb.set_model(self.stage1_discriminator)
        tb.set_model(self.ca_network)
        tb.set_model(self.embedding_compressor)

    def train_stage1(self):
        """Trains the stage1 StackGAN.
        """
        x_train, y_train, train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
                                        dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(64, 64))

        x_test, y_test, test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
      dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(64, 64))

        real = np.ones((self.batch_size, 1), dtype='float') * 0.9
        fake = np.zeros((self.batch_size, 1), dtype='float') * 0.1

        for epoch in range(self.epochs):
            print(f'Epoch: {epoch}')

            gen_loss = []
            dis_loss = []

            num_batches = int(x_train.shape[0] / self.batch_size)

            for i in range(num_batches):

                latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
                embedding_text = train_embeds[i * self.batch_size:(i + 1) * self.batch_size]
                compressed_embedding = self.embedding_compressor.predict_on_batch(embedding_text)
                compressed_embedding = np.reshape(compressed_embedding, (-1, 1, 1, 128))
                compressed_embedding = np.tile(compressed_embedding, (1, 4, 4, 1))

            image_batch = x_train[i * self.batch_size:(i+1) * self.batch_size]
            image_batch = (image_batch - 127.5) / 127.5

            gen_images, _ = self.stage1_generator.predict([embedding_text, latent_space])

            discriminator_loss = self.stage1_discriminator.train_on_batch([image_batch, compressed_embedding], 
                    np.reshape(real, (self.batch_size, 1)))

            discriminator_loss_gen = self.stage1_discriminator.train_on_batch([gen_images, compressed_embedding],
                    np.reshape(fake, (self.batch_size, 1)))

            discriminator_loss_wrong = self.stage1_discriminator.train_on_batch([gen_images[: self.batch_size-1], compressed_embedding[1:]], 
                    np.reshape(fake[1:], (self.batch_size-1, 1)))

# Discriminator loss
            d_loss = 0.5 * np.add(discriminator_loss, 0.5 * np.add(discriminator_loss_gen, discriminator_loss_wrong))
            dis_loss.append(d_loss)

            print(f'Discriminator Loss: {d_loss}')

               # Generator loss
            g_loss = self.stage1_adversarial.train_on_batch([embedding_text, latent_space, compressed_embedding],
                    [K.ones((self.batch_size, 1)) * 0.9, K.ones((self.batch_size, 256)) * 0.9])

            print(f'Generator Loss: {g_loss}')
            gen_loss.append(g_loss)

            if epoch % 5 == 0:
                    latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
                    embedding_batch = test_embeds[0 : self.batch_size]
                    gen_images, _ = self.stage1_generator.predict_on_batch([embedding_batch, latent_space])

                    for i, image in enumerate(gen_images[:10]):
                        save_image(image, f'test/gen_1_{epoch}_{i}')

            if epoch % 25 == 0:
                self.stage1_generator.save_weights('weights/stage1_gen.h5')
                self.stage1_discriminator.save_weights("weights/stage1_disc.h5")
                self.ca_network.save_weights('weights/stage1_ca.h5')
                self.embedding_compressor.save_weights('weights/stage1_embco.h5')
                self.stage1_adversarial.save_weights('weights/stage1_adv.h5')      

        self.stage1_generator.save_weights('weights/stage1_gen.h5')
        self.stage1_discriminator.save_weights("weights/stage1_disc.h5")


# In[ ]:


stage1 = StackGanStage1()
stage1.train_stage1()


# ### Check test folder for gernerated images from Stage1 Generator
# ### Let's Implement Stage 2 Generator

# In[29]:


############################################################
# Stage 2 Generator Network
############################################################

def concat_along_dims(inputs):
    """Joins the conditioned text with the encoded image along the dimensions.

    Args:
        inputs: consisting of conditioned text and encoded images as [c,x].

    Returns:
        Joint block along the dimensions.
    """
    c = inputs[0]
    x = inputs[1]

    c = K.expand_dims(c, axis=1)
    c = K.expand_dims(c, axis=1)
    c = K.tile(c, [1, 16, 16, 1])
    return K.concatenate([c, x], axis = 3)

def residual_block(input):
    """Residual block with plain identity connections.

    Args:
        inputs: input layer or an encoded layer

    Returns:
        Layer with computed identity mapping.
    """
    x = Conv2D(512, kernel_size=(3,3), padding='same', use_bias=False,
                kernel_initializer='he_uniform')(input)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    x = ReLU()(x)
    
    x = Conv2D(512, kernel_size=(3,3), padding='same', use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    
    x = add([x, input])
    x = ReLU()(x)

    return x

def build_stage2_generator():
    """Build the Stage 2 Generator Network using the conditioning text and images from stage 1.

    Returns:
        Stage 2 Generator Model for StackGAN.
    """
    input_layer1 = Input(shape=(1024,))
    input_images = Input(shape=(64, 64, 3))

    # Conditioning Augmentation
    ca = Dense(256)(input_layer1)
    mls = LeakyReLU(alpha=0.2)(ca)
    c = Lambda(conditioning_augmentation)(mls)

    # Downsampling block
    x = ZeroPadding2D(padding=(1,1))(input_images)
    x = Conv2D(128, kernel_size=(3,3), strides=1, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = ReLU()(x)

    x = ZeroPadding2D(padding=(1,1))(x)
    x = Conv2D(256, kernel_size=(4,4), strides=2, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    x = ReLU()(x)

    x = ZeroPadding2D(padding=(1,1))(x)
    x = Conv2D(512, kernel_size=(4,4), strides=2, use_bias=False,
                kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    x = ReLU()(x)

    # Concatenate text conditioning block with the encoded image
    concat = concat_along_dims([c, x])

    # Residual Blocks
    x = ZeroPadding2D(padding=(1,1))(concat)
    x = Conv2D(512, kernel_size=(3,3), use_bias=False, kernel_initializer='he_uniform')(x)
    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
    x = ReLU()(x)

    x = residual_block(x)
    x = residual_block(x)
    x = residual_block(x)
    x = residual_block(x)
    
    # Upsampling Blocks
    x = UpSamplingBlock(x, 512)
    x = UpSamplingBlock(x, 256)
    x = UpSamplingBlock(x, 128)
    x = UpSamplingBlock(x, 64)

    x = Conv2D(3, kernel_size=(3,3), padding='same', use_bias=False, kernel_initializer='he_uniform')(x)
    x = Activation('tanh')(x)
    
    stage2_gen = Model(inputs=[input_layer1, input_images], outputs=[x, mls])
    return stage2_gen


# In[30]:


generator_stage2 = build_stage2_generator()
generator_stage2.summary()


# In[31]:


############################################################
# Stage 2 Discriminator Network
############################################################

def build_stage2_discriminator():
    """Builds the Stage 2 Discriminator that uses the 256x256 resolution images from the generator
    and the compressed and spatially replicated embeddings.

    Returns:
        Stage 2 Discriminator Model for StackGAN.
    """
    input_layer1 = Input(shape=(256, 256, 3))

    x = Conv2D(64, kernel_size=(4,4), padding='same', strides=2, use_bias=False,
                kernel_initializer='he_uniform')(input_layer1)
    x = LeakyReLU(alpha=0.2)(x)

    x = ConvBlock(x, 128)
    x = ConvBlock(x, 256)
    x = ConvBlock(x, 512)
    x = ConvBlock(x, 1024)
    x = ConvBlock(x, 2048)
    x = ConvBlock(x, 1024, (1,1), 1)
    x = ConvBlock(x, 512, (1,1), 1, False)

    x1 = ConvBlock(x, 128, (1,1), 1)
    x1 = ConvBlock(x1, 128, (3,3), 1)
    x1 = ConvBlock(x1, 512, (3,3), 1, False)

    x2 = add([x, x1])
    x2 = LeakyReLU(alpha=0.2)(x2)

    # Concatenate compressed and spatially replicated embedding
    input_layer2 = Input(shape=(4, 4, 128))
    concat = concatenate([x2, input_layer2])

    x3 = Conv2D(512, kernel_size=(1,1), strides=1, padding='same', kernel_initializer='he_uniform')(concat)
    x3 = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x3)
    x3 = LeakyReLU(alpha=0.2)(x3)

    # Flatten and add a FC layer
    x3 = Flatten()(x3)
    x3 = Dense(1)(x3)
    x3 = Activation('sigmoid')(x3)

    stage2_dis = Model(inputs=[input_layer1, input_layer2], outputs=[x3])
    return stage2_dis


# In[32]:


discriminator_stage2 = build_stage2_discriminator()
discriminator_stage2.summary()


# In[33]:


############################################################
# Stage 2 Adversarial Model
############################################################

def stage2_adversarial_network(stage2_disc, stage2_gen, stage1_gen):
    """Stage 2 Adversarial Network.

    Args:
        stage2_disc: Stage 2 Discriminator Model.
        stage2_gen: Stage 2 Generator Model.
        stage1_gen: Stage 1 Generator Model.

    Returns:
        Stage 2 Adversarial network.
    """
    conditioned_embedding = Input(shape=(1024, ))
    latent_space = Input(shape=(100, ))
    compressed_replicated = Input(shape=(4, 4, 128))
    
    #the discriminator is trained separately and stage1_gen already trained, and this is the reason why we freeze its layers by setting the property trainable=false
    input_images, ca = stage1_gen([conditioned_embedding, latent_space])
    stage2_disc.trainable = False
    stage1_gen.trainable = False

    images, ca2 = stage2_gen([conditioned_embedding, input_images])
    probability = stage2_disc([images, compressed_replicated])

    return Model(inputs=[conditioned_embedding, latent_space, compressed_replicated],
        outputs=[probability, ca2])


# In[34]:


adversarial_stage2 = stage2_adversarial_network(discriminator_stage2, generator_stage2, generator)
adversarial_stage2.summary()


# In[35]:


class StackGanStage2(object):
    """StackGAN Stage 2 class.

    Args:
        epochs: Number of epochs
        z_dim: Latent space dimensions
        batch_size: Batch Size
        enable_function: If True, training function is decorated with tf.function
        stage2_generator_lr: Learning rate for stage 2 generator
        stage2_discriminator_lr: Learning rate for stage 2 discriminator
    """
    def __init__(self, epochs=500, z_dim=100, batch_size=64, enable_function=True, stage2_generator_lr=0.0002, stage2_discriminator_lr=0.0002):
        self.epochs = epochs
        self.z_dim = z_dim
        self.enable_function = enable_function
        self.stage1_generator_lr = stage2_generator_lr
        self.stage1_discriminator_lr = stage2_discriminator_lr
        self.low_image_size = 64
        self.high_image_size = 256
        self.conditioning_dim = 128
        self.batch_size = batch_size
        self.stage2_generator_optimizer = Adam(lr=stage2_generator_lr, beta_1=0.5, beta_2=0.999)
        self.stage2_discriminator_optimizer = Adam(lr=stage2_discriminator_lr, beta_1=0.5, beta_2=0.999)
        self.stage1_generator = build_stage1_generator()
        self.stage1_generator.compile(loss='binary_crossentropy', optimizer=self.stage2_generator_optimizer)
        self.stage1_generator.load_weights('weights/stage1_gen.h5')
        self.stage2_generator = build_stage2_generator()
        self.stage2_generator.compile(loss='binary_crossentropy', optimizer=self.stage2_generator_optimizer)

        self.stage2_discriminator = build_stage2_discriminator()
        self.stage2_discriminator.compile(loss='binary_crossentropy', optimizer=self.stage2_discriminator_optimizer)

        self.ca_network = build_ca_network()
        self.ca_network.compile(loss='binary_crossentropy', optimizer='Adam')

        self.embedding_compressor = build_embedding_compressor()
        self.embedding_compressor.compile(loss='binary_crossentropy', optimizer='Adam')

        self.stage2_adversarial = stage2_adversarial_network(self.stage2_discriminator, self.stage2_generator, self.stage1_generator)
        self.stage2_adversarial.compile(loss=['binary_crossentropy', adversarial_loss], loss_weights=[1, 2.0], optimizer=self.stage2_generator_optimizer)	

        self.checkpoint2 = tf.train.Checkpoint(
            generator_optimizer=self.stage2_generator_optimizer,
            discriminator_optimizer=self.stage2_discriminator_optimizer,
            generator=self.stage2_generator,
            discriminator=self.stage2_discriminator,
            generator1=self.stage1_generator)

    def visualize_stage2(self):
        """Running Tensorboard visualizations.
        """
        tb = TensorBoard(log_dir="logs/".format(time.time()))
        tb.set_model(self.stage2_generator)
        tb.set_model(self.stage2_discriminator)

    def train_stage2(self):
        """Trains Stage 2 StackGAN.
        """
        x_high_train, y_high_train, high_train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
            dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(256, 256))

        x_high_test, y_high_test, high_test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
              dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(256, 256))

        x_low_train, y_low_train, low_train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
          dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(64, 64))

        x_low_test, y_low_test, low_test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
          dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(64, 64))

        real = np.ones((self.batch_size, 1), dtype='float') * 0.9
        fake = np.zeros((self.batch_size, 1), dtype='float') * 0.1

        for epoch in range(self.epochs):
            print(f'Epoch: {epoch}')

            gen_loss = []
            disc_loss = []

            num_batches = int(x_high_train.shape[0] / self.batch_size)

            for i in range(num_batches):

                latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
                embedding_text = high_train_embeds[i * self.batch_size:(i + 1) * self.batch_size]
                compressed_embedding = self.embedding_compressor.predict_on_batch(embedding_text)
                compressed_embedding = np.reshape(compressed_embedding, (-1, 1, 1, self.conditioning_dim))
                compressed_embedding = np.tile(compressed_embedding, (1, 4, 4, 1))

                image_batch = x_high_train[i * self.batch_size:(i+1) * self.batch_size]
                image_batch = (image_batch - 127.5) / 127.5
                
                low_res_fakes, _ = self.stage1_generator.predict([embedding_text, latent_space], verbose=3)
                high_res_fakes, _ = self.stage2_generator.predict([embedding_text, low_res_fakes], verbose=3)

                discriminator_loss = self.stage2_discriminator.train_on_batch([image_batch, compressed_embedding],
                    np.reshape(real, (self.batch_size, 1)))

                discriminator_loss_gen = self.stage2_discriminator.train_on_batch([high_res_fakes, compressed_embedding],
                    np.reshape(fake, (self.batch_size, 1)))

                discriminator_loss_fake = self.stage2_discriminator.train_on_batch([image_batch[:(self.batch_size-1)], compressed_embedding[1:]],
                    np.reshape(fake[1:], (self.batch_size - 1, 1)))

                d_loss = 0.5 * np.add(discriminator_loss, 0.5 * np.add(discriminator_loss_gen, discriminator_loss_fake))
                disc_loss.append(d_loss)

                print(f'Discriminator Loss: {d_loss}')

                g_loss = self.stage2_adversarial.train_on_batch([embedding_text, latent_space, compressed_embedding],
                    [K.ones((self.batch_size, 1)) * 0.9, K.ones((self.batch_size, 256)) * 0.9])
                gen_loss.append(g_loss)

                print(f'Generator Loss: {g_loss}')

                if epoch % 5 == 0:
                    latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
                    embedding_batch = high_test_embeds[0 : self.batch_size]

                    low_fake_images, _ = self.stage1_generator.predict([embedding_batch, latent_space], verbose=3)
                    high_fake_images, _ = self.stage2_generator.predict([embedding_batch, low_fake_images], verbose=3)

                    for i, image in enumerate(high_fake_images[:10]):
                        save_image(image, f'results_stage2/gen_{epoch}_{i}.png')

                if epoch % 10 == 0:
                    self.stage2_generator.save_weights('weights/stage2_gen.h5')
                    self.stage2_discriminator.save_weights("weights/stage2_disc.h5")
                    self.ca_network.save_weights('weights/stage2_ca.h5')
                    self.embedding_compressor.save_weights('weights/stage2_embco.h5')
                    self.stage2_adversarial.save_weights('weights/stage2_adv.h5')

        self.stage2_generator.save_weights('weights/stage2_gen.h5')
        self.stage2_discriminator.save_weights("weights/stage2_disc.h5")


# In[ ]:


stage2 = StackGanStage2()
stage2.train_stage2()