#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun May  9 11:01:27 2021

@author: valentinahrtonova
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 30 23:05:23 2021

@author: valentinahrtonova
"""
#from __future__ import print_function, division

from keras.initializers import RandomNormal
from keras.layers import Input, Dropout, Concatenate, BatchNormalization, Activation
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.models import Model
from keras.optimizers import Adam

def build_generator(self):
        """U-Net Generator (G)"""

        def encoder_unet(layer_in, n_filters, batch_norm=True):
            """Downsampling with Conv2D"""
    
            e = Conv2D(filters = n_filters, kernel_size = (4,4), strides = (2,2), padding = 'same', kernel_initializer = self.init)(layer_in)
            if batch_norm:
                e = BatchNormalization(momentum = 0.8)(e)
            e = LeakyReLU(alpha = 0.2)(e)
    
            return e
    
        def decoder_unet(layer_in, skip_in, n_filters, dropout = True):
            """Upsampling with Conv2DTranspose"""
    
            d = Conv2DTranspose(filters = n_filters, kernel_size = (4,4), strides = (2,2), padding ='same', kernel_initializer = self.init)(layer_in)
            d = BatchNormalization(momentum = 0.8)(d)
            d = Activation('relu')(d)
            if dropout:
                d = Dropout(rate = 0.5)(d)
    
            # concatenate with skip connection
            d = Concatenate()([d, skip_in])
    
            return d
    
        gf = self.gf
    
        # Input image
        e0 = Input(shape=self.img_shape)

    # Encoder model
            # Encoder model
        e1 = encoder_unet(e0, gf, batch_norm=False)
        e2 = encoder_unet(e1, gf * 2)
        e3 = encoder_unet(e2, gf * 4)
        e4 = encoder_unet(e3, gf * 8)
        e5 = encoder_unet(e4, gf * 8)
        e6 = encoder_unet(e5, gf * 8)
        e7 = encoder_unet(e6, gf * 8)
    
    # Bottleneck    
        b = Conv2D(filters = 512, kernel_size = (4,4), strides = (2,2), padding = 'same', kernel_initializer = self.init)(e7)
        b = Activation('relu')(b)

    # Decoder model
        d1 = decoder_unet(b, e7, gf * 8)
        d2 = decoder_unet(d1, e6, gf * 8)
        d3 = decoder_unet(d2, e5, gf * 8)
        d4 = decoder_unet(d3, e4, gf * 8, dropout = False)
        d5 = decoder_unet(d4, e3, gf * 4, dropout = False)
        d6 = decoder_unet(d5, e2, gf * 2, dropout = False)
        d7 = decoder_unet(d6, e1, gf, dropout = False)
        
        # Output image
        g = Conv2DTranspose(filters = self.channels, kernel_size = (4,4), strides = (2,2), padding='same', kernel_initializer = self.init)(d7)
        image_out = Activation('tanh')(g)
    
        # Define model
        model = Model(e0, image_out)
        model.summary()
        
        return model

def build_discriminator(self):
        """ PatchGAN Discriminator (D)"""
    
        def discriminator_layer(layer_in, n_filters, n_stride = (2,2), batch_norm = True):
            """Discriminator layer"""
    
            d = Conv2D(filters = n_filters, kernel_size = (4,4) , strides = n_stride, padding = 'same', kernel_initializer = self.init)(layer_in)
            if batch_norm:
                d = BatchNormalization(momentum = 0.8)(d)
            
            d = LeakyReLU(alpha = 0.2)(d)
            return d
        
        df = self.df
        
        # Input images
        img_A = Input(shape=self.img_shape)
        img_B = Input(shape=self.img_shape)
    
        # Concatenate target image and conditioning image to produce input
        imgs_concat = Concatenate()([img_A, img_B])
    
        d1 = discriminator_layer(imgs_concat, df, batch_norm = False)
        d2 = discriminator_layer(d1, df * 2)
        d3 = discriminator_layer(d2, df * 4)
        d4 = discriminator_layer(d3, df * 8)
        d5 = discriminator_layer(d4, df * 8, n_stride = (1,1))
        
        # Patch output
        d6 = Conv2D(filters = 1, kernel_size = (4,4), padding = 'same', kernel_initializer = self.init)(d5)
        image_out = Activation('sigmoid')(d6)
    
        # Compile model
        model = Model([img_A, img_B], image_out)
        model.compile(loss='binary_crossentropy', optimizer=self.opt, loss_weights=[0.5])
        model.summary()
        
        return model

def build_gan(self):
    
    # Target images and their conditioning images
    img_A = Input(shape=self.img_shape)
    img_B = Input(shape=self.img_shape)

    # Define weight initializer
    self.init = RandomNormal(stddev=0.02)
    
    # Number of filters in G and D
    self.gf = 64
    self.df = 64
    
    # Calculate receptive field of D
    patch = int(self.img_rows / 2**4)
    self.disc_patch = (patch, patch, 1)

    # Define optimizer
    self.opt = Adam(learning_rate = 0.0002, beta_1 = 0.5)
    
    # Build G and D models
    #self.discriminator = load(self.directory + '/d_model_040.h5')
    #self.generator = load(self.directory + '/g_model_040.h5')
    self.discriminator = build_discriminator(self)
    self.generator = build_generator(self)

    # Generate a fake image based on image A (image B is target)
    fake_B = self.generator(img_A)
    
    # Initiate losses
    self.d_loss_real = list()
    self.d_loss_fake = list()
    self.d_loss = list()
    self.g_loss = list()

    # Only train G, set D weights to not trainable
    self.discriminator.trainable = False

    # Evaluate validity of generated image
    valid = self.discriminator([img_A, fake_B])
    
    # Compile GAN model
    #self.combined = load(self.directory + '/gan_model_040.h5')
    self.combined = Model(inputs = [img_A, img_B], outputs = [valid, fake_B])
    self.combined.compile(loss = ['binary_crossentropy', 'mae'], optimizer = self.opt, loss_weights = [1, 100])