#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 26 16:28:20 2021

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

from keras.layers import Dense, Reshape, Flatten, Dropout, BatchNormalization
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.models import Sequential

from keras.losses import BinaryCrossentropy

def build_gan(self):
    
    def build_generator(self):

        model = Sequential()
        
        gf = self.gf
        
        input_nodes = 256 * 4 * 4
        
        model.add(Dense(input_nodes, input_dim=self.latent_dim))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Reshape((4, 4, 256))) 
    	# upsample to 8x8
        model.add(Conv2DTranspose(gf*4, (4,4), strides= (2,2), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
    	# upsample to 16x16
        model.add(Conv2DTranspose(gf*4, (4,4), strides= (2,2), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
    	# upsample to 32x32
        model.add(Conv2DTranspose(gf*2, (4,4), strides= (2,2), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        # upsample to 64x64
        model.add(Conv2DTranspose(gf*2, (4,4), strides= (2,2), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))	
        # upsample to 128x128
        model.add(Conv2DTranspose(gf, (4,4), strides= (2,2), padding='same', kernel_initializer=self.init))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        # upsample to 256x256
    	# output layer
        model.add(Conv2D(3, (3,3), activation='tanh', padding='same'))
    
        model.summary()
    
        return model

    def build_discriminator(self):
        
        df = self.df
    
        model = Sequential()
    
        model.add(Conv2D(df, (4,4), padding='same', input_shape=self.img_shape, kernel_initializer=self.init))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.4))
        # downsample to 64x64
        model.add(Conv2D(df*2, (4,4), 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.4))
    	# downsample 32x32
        model.add(Conv2D(df*2, (4,4), 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.4))
    	# downsample 16x16
        model.add(Conv2D(df*4, (4,4), 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.4))
    	# downsample 8x8
        model.add(Conv2D(df*4, (4,4), 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.4))
        # downsample 4x4
        model.add(Conv2D(df*8, (4,4), 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.4))
    	# classifier
        model.add(Flatten())
        model.add(Dropout(0.4))
        model.add(Dense(1, activation='sigmoid'))
    	# compile model
        model.compile(loss=BinaryCrossentropy(label_smoothing = 0.1), optimizer=self.opt)
        model.summary()
        
        return model
    
    self.discriminator = build_discriminator(self)
    self.generator = build_generator(self)
    
    # Only train generator for combined model
    self.discriminator.trainable = False
    
    model = Sequential()
    
    model.add(self.generator)
    model.add(self.discriminator)
    
    model.compile(loss=BinaryCrossentropy(label_smoothing = 0.1), optimizer=self.opt)
    return model

