import tensorflow as tf
import tensorflow.contrib.slim as slim

from utils import lrelu

import numpy as np

class Model(object):
    
    """
    Reference: Volpi et al., Adversarial Feature Augmentation 
    for Unsupervised Domain Adaptation, CVPR 2018
    
    Author: Riccardo Volpi
    """
    
    def __init__(self, mode='train_feature_generator', learning_rate=0.0003):
        
	self.mode = mode
        self.learning_rate = learning_rate
	self.hidden_repr_size = 128
    
    def feature_extractor(self, images, reuse=False, return_output = False):
	
	'''
	Takes in input images and gives in output:	
	1. logits, in the pretrain phase
	2. features, in the feature generation phase 
	'''
	
	# for SVHN images
	if images.get_shape()[3] == 3:
	    images = tf.image.rgb_to_grayscale(images)

	
	with tf.variable_scope('feature_extractor', reuse=reuse):
	    with slim.arg_scope([slim.conv2d], activation_fn=tf.nn.relu, padding='VALID'):
		net = slim.conv2d(images, 64, 5, scope='conv1')
		net = slim.max_pool2d(net, 2, stride=2, scope='pool1')
		net = slim.conv2d(net, 128, 5, scope='conv2')
		net = slim.max_pool2d(net, 2, stride=2, scope='pool2')
		net = tf.contrib.layers.flatten(net)
		net = slim.fully_connected(net, 1024, activation_fn=tf.nn.relu, scope='fc3')
		net = slim.fully_connected(net, self.hidden_repr_size, activation_fn=tf.tanh, scope='fc4') # has to be tanh, as it's more welcomed by GANs
		
		if (self.mode == 'train_feature_extractor') or return_output: # in the pretraining phase we need logits, in the feature generation phase we need features
		    net = slim.fully_connected(net, 10, activation_fn=None, scope='fc5')
		
		return net

    def feature_generator(self, noise, labels, reuse=False):
	
	'''
	Takes in input noise and labels, and
	generates f(z|y)  
	'''
	
	try: # just to make it work on different Tensorflow releases
	    inputs = tf.concat(1, [noise, tf.cast(labels,tf.float32)])
	except:
	    inputs = tf.concat([noise, tf.cast(labels,tf.float32)], 1)
	
	with tf.variable_scope('feature_generator', reuse=reuse):
	    with slim.arg_scope([slim.fully_connected], weights_initializer=tf.contrib.layers.xavier_initializer(), biases_initializer = tf.constant_initializer(0.0)):
		
		with slim.arg_scope([slim.batch_norm], decay=0.95, center=True, scale=True, 
                                    activation_fn=tf.nn.relu, is_training=(self.mode=='train_feature_generator')):
                    
		    net = slim.fully_connected(inputs, 1024, activation_fn = tf.nn.relu, scope='sgen_fc1')
		    net = slim.batch_norm(net, scope='sgen_bn1')
		    net = slim.dropout(net, 0.5) #dropout needs to be always on, do not include "is_training"
		    net = slim.fully_connected(net, 1024, activation_fn = tf.nn.relu, scope='sgen_fc2')
		    net = slim.batch_norm(net, scope='sgen_bn2')
		    net = slim.dropout(net, 0.5)
		    net = slim.fully_connected(net, self.hidden_repr_size, activation_fn = tf.tanh, scope='sgen_feat')
		    return net
				
    def feature_discriminator(self, features, labels, reuse=False):
	
	'''
	Used in Step 1 to discriminate between 
	real and generated features. Takes in 
	input features and associated labels and 
	has to discriminate whether they are real 
	or generated by feature_generator  
	'''
		
	try: # just to make it work on different Tensorflow releases
	    inputs = tf.concat(1, [features, tf.cast(labels,tf.float32)])
	except:
	    inputs = tf.concat([features, tf.cast(labels,tf.float32)], 1)
	
	with tf.variable_scope('feature_discriminator',reuse=reuse):
	    with slim.arg_scope([slim.fully_connected],weights_initializer=tf.contrib.layers.xavier_initializer(), biases_initializer = tf.constant_initializer(0.0)):		
		net = slim.fully_connected(inputs, 1024, activation_fn = tf.nn.relu, scope='sdisc_fc1')
		net = slim.fully_connected(net,1,activation_fn=tf.sigmoid,scope='sdisc_prob')
		return net
				
    def feature_discriminator_DIFA(self, features, reuse=False):
	
	'''
	Used in Step 2 to discriminate between 
	source and target features. Takes in input 
	features and has to discriminate whether 
	they are from S or from the encoder we are 
	training.  
	'''

	with tf.variable_scope('feature_discriminator_DIFA',reuse=reuse):
	    with slim.arg_scope([slim.fully_connected],weights_initializer=tf.contrib.layers.xavier_initializer(), biases_initializer = tf.constant_initializer(0.0)):		
		net = slim.fully_connected(features, 1024, activation_fn = lrelu, scope='sdisc_DIFA_fc1')
		net = slim.fully_connected(net,1,activation_fn=tf.sigmoid,scope='sdisc_DIFA_prob')
		return net

    def build_model(self):
        
        if self.mode == 'train_feature_extractor':
	    
	    #################################################################################################################
	    # Step 0: training a classifier to classify well some images (source images, if using it for domain adaptation) #
	    #################################################################################################################
	    
	    
	    self.images = tf.placeholder(tf.float32, [None, 32, 32, 3], 'images')
	    self.labels = tf.placeholder(tf.int64, [None], 'labels')
            
	    # extracting logits through feature_extractor (C(E_S) in the paper)
	    self.logits = self.feature_extractor(self.images)
	    
	    	
	    self.pred = tf.argmax(self.logits, 1)
            self.correct_pred = tf.equal(self.pred, self.labels)
            self.accuracy = tf.reduce_mean(tf.cast(self.correct_pred, tf.float32))
            
            self.loss = slim.losses.sparse_softmax_cross_entropy(self.logits, self.labels)
            self.optimizer = tf.train.AdamOptimizer(self.learning_rate) 
            self.train_op = slim.learning.create_train_op(self.loss, self.optimizer)
	    
            # summary
            loss_summary = tf.summary.scalar('classification_loss', self.loss)
            src_accuracy_summary = tf.summary.scalar('accuracy', self.accuracy)
            self.summary_op = tf.summary.merge([loss_summary, src_accuracy_summary])
	
	elif self.mode == 'train_feature_generator':
	    
	    #######################################################################################################################
	    # Step 1: training a feature generator to generate realistic features conditioning on a noise vector and a label code #
	    #######################################################################################################################
	    
				
	    self.images = tf.placeholder(tf.float32, [None, 32, 32, 3], 'images')
	    self.noise = tf.placeholder(tf.float32, [None, 100], 'noise')
	    self.labels = tf.placeholder(tf.int64, [None, 10], 'labels')
	    
	    # extracting features associated to our images through extractor (E_S in the paper)
	    self.real_features = self.feature_extractor(self.images)
	    
	    # features generated by feature_generator (S in the paper)
	    self.gen_features = self.feature_generator(self.noise, self.labels) 

	    # In the following:
	    #
	    # 'real': related to features extracted from images.
	    # 'fake': related to generated features.
	    # The minimax game is solved in the Least Squares GAN framework.
	    # Ones are related to real features, zeros to generated features. Generator is trained by flipping labels. 

	    self.logits_real = self.feature_discriminator(self.real_features,self.labels) 
	    self.logits_fake = self.feature_discriminator(self.gen_features,self.labels, reuse=True)
	    
	    self.d_loss_real = tf.reduce_mean(tf.square(self.logits_real - tf.ones_like(self.logits_real)))  
	    self.d_loss_fake = tf.reduce_mean(tf.square(self.logits_fake - tf.zeros_like(self.logits_fake))) 
	    
	    self.d_loss = self.d_loss_real + self.d_loss_fake
	    
	    self.g_loss = tf.reduce_mean(tf.square(self.logits_fake - tf.ones_like(self.logits_fake))) # flipping labels        
	    
	    self.d_optimizer = tf.train.AdamOptimizer(self.learning_rate)
	    self.g_optimizer = tf.train.AdamOptimizer(self.learning_rate)
	    
	    t_vars = tf.trainable_variables()
	    d_vars = [var for var in t_vars if 'feature_discriminator' in var.name]
	    g_vars = [var for var in t_vars if 'feature_generator' in var.name]
	    
	    
	    with tf.variable_scope('train_op',reuse=False):
		
		# training feature_discriminator to classify real features as ones and fake features as zeros.
		self.d_train_op = slim.learning.create_train_op(self.d_loss, self.d_optimizer, variables_to_train=d_vars) 
		
		# training feature_generator to make feature_discriminator classify fake features as ones, i.e. to generate more realistic features.
		self.g_train_op = slim.learning.create_train_op(self.g_loss, self.g_optimizer, variables_to_train=g_vars) 
	    
	    # summary 
	    d_loss_summary = tf.summary.scalar('d_loss', self.d_loss)
	    g_loss_summary = tf.summary.scalar('g_loss', self.g_loss)
	    self.summary_op = tf.summary.merge([d_loss_summary, g_loss_summary])

	    for var in tf.trainable_variables():
		tf.summary.histogram(var.op.name, var)
       
        elif self.mode == 'train_DIFA':
	    
	    ########################################################################################################
	    # Step 2: training a feature extractor to map both source and target samples in the same feature space #
	    ########################################################################################################

	    self.noise = tf.placeholder(tf.float32, [None, 100], 'noise')
            self.labels = tf.placeholder(tf.float32, [None, 10], 'labels') # to generate features with S 
            self.src_images = tf.placeholder(tf.float32, [None, 32, 32, 3], 'svhn_images')
            self.trg_images = tf.placeholder(tf.float32, [None, 32, 32, 1], 'mnist_images')
	    
	    self.trg_labels = self.feature_extractor(self.trg_images, return_output=True)
	    
	    #For testing.......................................................................
	    self.trg_labels_gt = tf.placeholder(tf.int64, [None], 'mnist_labels_for_testing') #ground truth for target, used for testing
	    self.trg_pred = tf.argmax(self.trg_labels, 1)
	    self.trg_correct_pred = tf.equal(self.trg_pred, self.trg_labels_gt)
	    self.trg_accuracy = tf.reduce_mean(tf.cast(self.trg_correct_pred, tf.float32))
	    #..................................................................................

	    try: #try/except due to differences among tf versions
		self.images = tf.concat(0, [tf.image.rgb_to_grayscale(self.src_images), self.trg_images])
	    except:
		self.images = tf.concat([tf.image.rgb_to_grayscale(self.src_images), self.trg_images], 0)
		
	    self.features_from_S = self.feature_generator(self.noise,self.labels)		
	    self.features_from_E = self.feature_extractor(self.images, reuse=True)
	    
	    self.logits_real = self.feature_discriminator_DIFA(self.features_from_S)
	    self.logits_fake = self.feature_discriminator_DIFA(self.features_from_E, reuse=True)
	    
	    # Discriminator loss
	    self.d_loss_real = tf.reduce_mean(tf.square(self.logits_real - tf.ones_like(self.logits_real)))
	    self.d_loss_fake = tf.reduce_mean(tf.square(self.logits_fake - tf.zeros_like(self.logits_fake)))
	    self.d_loss = self.d_loss_real + self.d_loss_fake 
	    
	    # Encoder loss
	    self.e_loss = tf.reduce_mean(tf.square(self.logits_fake - tf.ones_like(self.logits_fake)))
	    
	    # Optimizers
	    self.d_optimizer = tf.train.AdamOptimizer(self.learning_rate)
            self.e_optimizer = tf.train.AdamOptimizer(self.learning_rate)
            
            t_vars = tf.trainable_variables()
            e_vars = [var for var in t_vars if 'feature_extractor' in var.name]
            d_vars = [var for var in t_vars if 'feature_discriminator_DIFA' in var.name]
            
            # train op
            with tf.variable_scope('training_op',reuse=False):
                self.e_train_op = slim.learning.create_train_op(self.e_loss, self.e_optimizer, variables_to_train=e_vars)
                self.d_train_op = slim.learning.create_train_op(self.d_loss, self.d_optimizer, variables_to_train=d_vars)
            
            # summary op
            e_loss_summary = tf.summary.scalar('G_loss', self.e_loss)
            d_loss_summary = tf.summary.scalar('D_loss', self.d_loss)
            self.summary_op = tf.summary.merge([e_loss_summary, d_loss_summary])
            

            for var in tf.trainable_variables():
		tf.summary.histogram(var.op.name, var) 
        
        else:
	    raise Exception('Unknown mode.')












