#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Author   : MuyunLi
Date     : 6/22/20 6:14 PM
FileName : train.py
Training DCGAN
"""
import time
from datetime import datetime
import os
import glob
import numpy as np
from scipy import misc

from network import *
import logging

logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                    level=logging.INFO)

def train():
    # Make sure that the images folder containing all images is in the same directory as all Python files
    if not os.path.exists("target_images"):
        raise Exception("The images folder containing all pictures is not in this directory, please add")

    # Get training data
    data = []
    for image in glob.glob("target_images/*"):
        image_data = misc.imread(image)  # imread uses PIL to read image data
        data.append(image_data)
    input_data = np.array(data)

    # Normalize the data to the value of [-1, 1], which is also the output range of the Tanh activation function
    input_data = (input_data.astype(np.float32) - 127.5) / 127.5

    # Construct generator and discriminator
    g = generator_model()
    d = discriminator_model()

    # Construct a network model composed of generator and discriminator
    d_on_g = generator_containing_discriminator(g, d)

    # Adam Optimizer
    g_optimizer = tf.keras.optimizers.Adam(lr=LEARNING_RATE, beta_1=BETA_1)
    d_optimizer = tf.keras.optimizers.Adam(lr=LEARNING_RATE, beta_1=BETA_1)

    # Configure generator and discriminator
    g.compile(loss="binary_crossentropy", optimizer=g_optimizer)
    d_on_g.compile(loss="binary_crossentropy", optimizer=g_optimizer)
    d.trainable = True
    d.compile(loss="binary_crossentropy", optimizer=d_optimizer)

    # training
    for epoch in range(EPOCHS):
        for index in range(int(input_data.shape[0] / BATCH_SIZE)):
            input_batch = input_data[index * BATCH_SIZE : (index + 1) * BATCH_SIZE]

            # Continuous uniformly distributed random data (noise)
            random_data = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100))
            # Image data generated by the generator
            generated_images = g.predict(random_data, verbose=0)

            input_batch = np.concatenate((input_batch, generated_images))
            output_batch = [1] * BATCH_SIZE + [0] * BATCH_SIZE

            # Train the discriminator to give it the ability to identify unqualified images
            d_loss = d.train_on_batch(input_batch, output_batch)

            # When training the generator, make the discriminator untrainable
            d.trainable = False

            # Regenerate random data. Very critical
            random_data = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100))

            # Train the generator and use a discriminator that cannot be trained to discriminate
            g_loss = d_on_g.train_on_batch(random_data, [1] * BATCH_SIZE)

            # Recovery discriminator can be trained
            d.trainable = True

            # Print loss
            logging.info("Epoch {}, step {}, generator loss: {:.3f}, discriminator loss: {:.3f}".format(epoch, index, g_loss, d_loss))

        # Save the parameters of the generator and discriminator
        if epoch % 10 == 9:
            g.save_weights("generator_weight", True)
            d.save_weights("discriminator_weight", True)




if __name__ == "__main__":
    train()

