#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun May 16 10:20:30 2021

@author: valentinahrtonova
"""

# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
from keras.layers import Input, Dense, Reshape, Flatten, Dropout
from keras.layers import BatchNormalization, Activation, ZeroPadding2D, AveragePooling2D
from keras.layers.advanced_activations import LeakyReLU, ReLU
from keras.layers.convolutional import UpSampling2D, Conv2D, Conv2DTranspose
from keras.models import Sequential, Model
from keras.optimizers import RMSprop
from keras.activations import linear
from keras.initializers import RandomNormal

from numpy.random import randn

import keras.backend as K

import matplotlib.pyplot as plt
import os
import sys

import numpy as np
from numpy import load

class WGAN():
    def __init__(self):
        self.img_rows = 64
        self.img_cols = 64
        self.channels = 3
        self.img_shape = (self.img_rows, self.img_cols, self.channels)
        self.latent_dim = 100
        
        self.directory = 'melanoma_WGAN_64_5'
        self.dataset = 'melanoma_dataset_full_64.npz'
        
        # Following parameter and optimizer set as recommended in paper
        self.n_critic = 5
        self.clip_value = 0.005
        self.opt = RMSprop(lr=0.00005)
        
        self.init = RandomNormal(mean=0.0, stddev=0.02, seed=None)

        # Build and compile the critic
        self.critic = self.build_critic()
        self.critic.compile(loss=self.wasserstein_loss,
            optimizer=self.opt,
            metrics=['accuracy'])
        
        # Build the generator
        self.generator = self.build_generator()

        # The generator takes noise as input and generated imgs
        z = Input(shape=(self.latent_dim,))
        img = self.generator(z)

        # For the combined model we will only train the generator
        self.critic.trainable = False

        # The critic takes generated images as input and determines validity
        valid = self.critic(img)

        # The combined model  (stacked generator and critic)
        self.combined = Model(z, valid)
        self.combined.compile(loss=self.wasserstein_loss,
            optimizer=self.opt,
            metrics=['accuracy'])

    def wasserstein_loss(self, y_true, y_pred):
        return K.mean(y_true * y_pred)

    def build_generator(self):
        
        n_filters = 32

        model = Sequential()

        n_nodes = n_filters*8 * 4 * 4
    
        model.add(Dense(n_nodes, input_dim=self.latent_dim))
        model.add(Activation("relu"))
        model.add(Reshape((4, 4, n_filters*8))) 
        model.add(UpSampling2D())
        #psample to 8x8
        model.add(Conv2D(n_filters*8, (4,4), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(UpSampling2D())
        #psample to 16x16
        model.add(Conv2D(n_filters*4, (4,4), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(UpSampling2D())
    # upsample to 64x64
        model.add(Conv2D(n_filters*2, (4,4), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(UpSampling2D())
        # upsample to 64x64
        model.add(Conv2D(n_filters, (4,4), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
	# output layer
        model.add(Conv2D(3, (4,4), activation='tanh', padding='same'))

        model.summary()

        noise = Input(shape=(self.latent_dim,))
        img = model(noise)

        return Model(noise, img)

    def build_critic(self):

        num_filters = 32
        model = Sequential()
	# normal
        model.add(Conv2D(num_filters, (3,3), padding='same', input_shape=self.img_shape, kernel_initializer=self.init))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
    # downsample
        model.add(Conv2D(num_filters*2, (3,3), strides=(2,2), padding='same', kernel_initializer=self.init))
        #model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
	# downsample
        model.add(Conv2D(num_filters*4, (3,3), strides=(2,2), padding='same', kernel_initializer=self.init))
        #model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
	# downsample
        model.add(Conv2D(num_filters*4, (3,3), strides=(2,2), padding='same', kernel_initializer=self.init))
        #model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
	# downsample
        model.add(Conv2D(num_filters*8, (3,3), strides=(2,2), padding='same', kernel_initializer=self.init))
       #model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
	# classifier
        model.add(Flatten())
        #model.add(Dropout(0.4))
        model.add(Dense(1))
	# compile model
        model.compile(loss = self.wasserstein_loss, optimizer = self.opt)
        model.summary()

        img = Input(shape=self.img_shape)
        validity = model(img)

        return Model(img, validity)
    # generate points in latent space as input for the generator
    def generate_latent_points(self, latent_dim, n_samples):
	# generate points in the latent space
    	x_input = randn(latent_dim * n_samples)
	# reshape into a batch of inputs for the network
    	x_input = x_input.reshape(n_samples, latent_dim)
    	return x_input

    def train(self, epochs, batch_size, sample_interval):
        
        # Load the dataset
        data = load(self.dataset)
    # extract numpy array
        X = data['arr_0']
        X.astype('float32')
	# scale from [0,255] to [-1,1]
        X_train = (X - 127.5) / 127.5
        
        bat_per_epo = int(X_train.shape[0] / batch_size)   
        half_batch = int(batch_size / 2)
        
        print('>Epochs: %s, Batch per epoch: %s' % (epochs, bat_per_epo))

        # Adversarial ground truths
        valid = -np.ones((half_batch, 1))
        fake = np.ones((half_batch, 1))

        for epoch in range(epochs):
            
            for i in range (bat_per_epo):

                for _ in range(self.n_critic):
    
                    # Select a random batch of images
                    idx = np.random.randint(0, X_train.shape[0], half_batch)
                    imgs = X_train[idx]
                    
                    # Sample noise as generator input
                    noise = self.generate_latent_points(self.latent_dim, half_batch)
    
                    # Generate a batch of new images
                    gen_imgs = self.generator.predict(noise)
    
                    # Train the critic
                    d_loss_real = self.critic.train_on_batch(imgs, valid)
                    d_loss_fake = self.critic.train_on_batch(gen_imgs, fake)
                    d_loss = 0.5 * np.add(d_loss_fake, d_loss_real)
    
                    # Clip critic weights
                    for l in self.critic.layers:
                        weights = l.get_weights()
                        weights = [np.clip(w, -self.clip_value, self.clip_value) for w in weights]
                        l.set_weights(weights)
    
                g_loss = self.combined.train_on_batch(noise, valid)
    
                # Plot the progress
                print ("Epoch: %d Batch: %d [D loss: %f] [G loss: %f]" % (epoch, i, 1 - d_loss[0], 1 - g_loss[0]))

            # If at save interval => save generated image samples
            if epoch % sample_interval == 0:
                self.sample_images(epoch)

    def sample_images(self, epoch):
        os.makedirs(self.directory, exist_ok=True)
        r, c = 5, 5
        noise = np.random.normal(0, 1, (r * c, self.latent_dim))
        gen_imgs = self.generator.predict(noise)

        # Rescale images 0 - 1
        gen_imgs = 0.5 * gen_imgs + 0.5

        fig, axs = plt.subplots(r, c)
        cnt = 0
        for i in range(r):
            for j in range(c):
                axs[i,j].imshow(gen_imgs[cnt])
                axs[i,j].axis('off')
                cnt += 1
        fig.savefig(self.directory + "/generated_%d.png" % epoch)
        g_filename = self.directory + '/g_model_%03d.h5' % (epoch+1)
        self.generator.save(g_filename)
        d_filename = self.directory + '/d_model_%03d.h5' % (epoch+1)
        self.critic.save(d_filename)
        gan_filename = self.directory + '/gan_model_%03d.h5' % (epoch+1)
        print('>Saved: %s' % (g_filename))
        self.combined.save(gan_filename)
        plt.close()

if __name__ == '__main__':
    wgan = WGAN()
    wgan.train(epochs=4000, batch_size=16, sample_interval=1)


