import os
import sys
import time
import skimage
import numpy as np
import scipy.io as sio
from tqdm import trange
import tensorflow as tf
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
from sklearn.utils import shuffle
from contextlib import redirect_stdout

from Tools import *
from Networks import *
from flip_gradient import flip_gradient
class Models():
    def __init__(self, args, dataset):
        self.args = args

        # Initializing the placeholders
        #Changing  the seed  at any run
        tf.set_random_seed(int(time.time()))
        tf.reset_default_graph()
        self.learning_rate = tf.placeholder(tf.float32, [], name="learning_rate")
        if self.args.compute_ndvi:
            self.data = tf.placeholder(tf.float32, [None, self.args.patches_dimension, self.args.patches_dimension, 2 * self.args.image_channels + 2], name = "data")
        else:
            self.data = tf.placeholder(tf.float32, [None, self.args.patches_dimension, self.args.patches_dimension, 2 * self.args.image_channels], name = "data")

        self.label = tf.placeholder(tf.float32, [None, self.args.num_classes], name = "label")
        self.label_d = tf.placeholder(tf.float32, [None, 2], name = "label_d")
        self.mask_c = tf.placeholder(tf.float32, [None,], name="labeled_samples")
        self.L = tf.placeholder(tf.float32, [], name="L" )

        # Initializing the network class
        self.classifier = EF_CNN(self.args)
        # Initializing the models individually
        if self.args.FE_Architecture == 'Mabel_Arch':
            Encoder_Outputs = self.classifier.build_Mabel_Arch(self.data, reuse = False, name = "FE")
        elif self.args.FE_Architecture == 'Ganin_Arch':
            Encoder_Outputs = self.classifier.build_Ganin_Arch(self.data, reuse = False, name = "FE")
        #Defining the classifier
        Classifier_Outputs = self.classifier.build_MLP_1hidden_cl(Encoder_Outputs[-1], reuse = False, name = "MLP_Cl")

        self.logits_c = Classifier_Outputs[-2]
        self.prediction_c = Classifier_Outputs[-1]
        self.features_c = Encoder_Outputs[-1]

        if self.args.training_type == 'domain_adaptation':
            if 'DR' in self.args.da_type:
                flip_feature = flip_gradient(self.features_c, self.L)
                self.DR = Domain_Regressors(self.args)
                DR_Ouputs = self.DR.build_Domain_Classifier_Arch(flip_feature, name = 'FC_Domain_Classifier')
                self.logits_d = DR_Ouputs[-2]

        if self.args.phase == 'train':
            self.summary(Encoder_Outputs, 'Encoder:')
            self.summary(Classifier_Outputs, 'Classifier:')

            self.dataset_s = dataset[0]
            self.dataset_t = dataset[1]
            #Defining losses
            temp = tf.nn.softmax_cross_entropy_with_logits(logits = self.logits_c, labels = self.label)
            self.classifier_loss =  tf.reduce_sum(self.mask_c * temp) / tf.reduce_sum(self.mask_c)
            if self.args.training_type == 'classification':
                self.total_loss = self.classifier_loss
            else:
                if 'DR' in self.args.da_type:
                    self.summary(DR_Ouputs, "Domain_Regressor: ")
                    self.domainregressor_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = self.logits_d, labels = self.label_d))
                    self.total_loss = self.classifier_loss + self.domainregressor_loss
                else:
                    self.total_loss = self.classifier_loss
            #Defining the Optimizers
            self.training_optimizer = tf.train.MomentumOptimizer(self.learning_rate, self.args.beta1).minimize(self.total_loss)
            self.saver = tf.train.Saver(max_to_keep=5)
            self.sess=tf.Session()
            self.sess.run(tf.initialize_all_variables())

        elif self.args.phase == 'test':
            self.dataset = dataset
            self.saver = tf.train.Saver(max_to_keep=5)
            self.sess=tf.Session()
            self.sess.run(tf.initialize_all_variables())
            print('[*]Loading the feature extractor and classifier trained models...')
            mod = self.load(self.args.trained_model_path)
            if mod:
                print(" [*] Load with SUCCESS")
            else:
                print(" [!] Load failed...")
                sys.exit()

    def Learning_rate_decay(self):
        lr = self.args.lr / (1. + 10 * self.p)**0.75
        return lr

    def summary(self, net, name):
        print(net)
        f = open(self.args.save_checkpoint_path + "Architecture.txt","a")
        f.write(name + "\n")
        for i in range(len(net)):
            print(net[i].get_shape().as_list())
            f.write(str(net[i].get_shape().as_list()) + "\n")
        f.close()

    def Train(self):

        pat = 0
        best_val_dr = 0
        best_val_fs = 0
        if self.args.balanced_tr:
            # Shuffling the data and labels
            central_pixels_coor_tr = self.dataset_s.central_pixels_coor_tr.copy()
            y_train = self.dataset_s.y_train.copy()
            central_pixels_coor_tr, y_train = shuffle(central_pixels_coor_tr, y_train, random_state=0)
            positive_coordinates = np.transpose(np.array(np.where(y_train == 1)))
            negative_coordinates = np.transpose(np.array(np.where(y_train == 0)))
            positive_coordinates = positive_coordinates[:,0]
            negative_coordinates = negative_coordinates[:,0]

            positive_central_pixels_coor_tr = central_pixels_coor_tr[positive_coordinates, :]
            if self.args.data_augmentation:
                positive_central_pixels_coor_tr, _ = Data_Augmentation_Definition(positive_central_pixels_coor_tr, np.ones((len(positive_coordinates),1)))

            #Taking the same amount of negative samples as positive
            negative_coordinates = negative_coordinates[:positive_central_pixels_coor_tr.shape[0]]
            if self.args.data_augmentation:
                negative_central_pixels_coor_tr = np.concatenate((central_pixels_coor_tr[negative_coordinates, :], np.zeros((len(negative_coordinates),1))),axis=1)
            else:
                negative_central_pixels_coor_tr = central_pixels_coor_tr[negative_coordinates, :]

            positive_y_train = np.ones((positive_central_pixels_coor_tr.shape[0],1))
            negative_y_train = np.zeros((negative_central_pixels_coor_tr.shape[0],1))
            central_pixels_coor_tr = np.concatenate((positive_central_pixels_coor_tr, negative_central_pixels_coor_tr), axis=0)
            y_train = np.concatenate((positive_y_train, negative_y_train), axis=0)
            # Shuffling again
            central_pixels_coor_tr_s, y_train_s = shuffle(central_pixels_coor_tr, y_train,random_state=0)
        if self.args.balanced_vl:
            central_pixels_coor_vl = self.dataset_s.central_pixels_coor_vl.copy()
            y_valid = self.dataset_s.y_valid.copy()
            # Shuffling the data and labels
            central_pixels_coor_vl, y_valid = shuffle(central_pixels_coor_vl, y_valid,random_state=0)

            positive_coordinates = np.transpose(np.array(np.where(y_valid == 1)))
            negative_coordinates = np.transpose(np.array(np.where(y_valid == 0)))
            positive_coordinates = positive_coordinates[:,0]
            negative_coordinates = negative_coordinates[:,0]

            positive_central_pixels_coor_vl = central_pixels_coor_vl[positive_coordinates, :]
            if self.args.data_augmentation:
                positive_central_pixels_coor_vl, _ = Data_Augmentation_Definition(positive_central_pixels_coor_vl, np.ones((len(positive_coordinates),1)))

            #Taking the same amount of negative samples as positive
            negative_coordinates = negative_coordinates[:positive_central_pixels_coor_vl.shape[0]]
            if self.args.data_augmentation:
                negative_central_pixels_coor_vl = np.concatenate((central_pixels_coor_vl[negative_coordinates, :] , np.zeros((len(negative_coordinates),1))), axis=1)
            else:
                negative_central_pixels_coor_vl = central_pixels_coor_vl[negative_coordinates, :]

            positive_y_valid = np.ones((positive_central_pixels_coor_vl.shape[0],1))
            negative_y_valid = np.zeros((negative_central_pixels_coor_vl.shape[0],1))
            central_pixels_coor_vl = np.concatenate((positive_central_pixels_coor_vl, negative_central_pixels_coor_vl), axis=0)
            y_valid = np.concatenate((positive_y_valid, negative_y_valid), axis=0)
            # Shuffling again
            central_pixels_coor_vl_s, y_valid_s = shuffle(central_pixels_coor_vl, y_valid,random_state=0)

        print('Source sets dimensions')
        print(np.shape(central_pixels_coor_tr_s))
        print(np.shape(central_pixels_coor_vl_s))
        print(np.shape(y_train_s))
        print(np.shape(y_valid_s))

        if self.args.training_type == 'classification':
            print('Classification training on source domain')
            y_train_s_hot = tf.keras.utils.to_categorical(y_train_s, self.args.num_classes)
            y_valid_s_hot = tf.keras.utils.to_categorical(y_valid_s, self.args.num_classes)

            central_pixels_coor_tr = central_pixels_coor_tr_s.copy()
            central_pixels_coor_vl = central_pixels_coor_vl_s.copy()

            y_train_c_hot = y_train_s_hot.copy()
            y_valid_c_hot = y_valid_s_hot.copy()
            y_train_d_hot = np.ones((y_train_c_hot.shape[0], 2))
            y_valid_d_hot = np.ones((y_valid_c_hot.shape[0], 2))
            classification_mask_tr = np.ones((y_train_c_hot.shape[0], 1))
            classification_mask_vl = np.ones((y_valid_c_hot.shape[0], 1))

            domain_indexs_tr = np.zeros((y_train_c_hot.shape[0], 1))
            domain_indexs_vl = np.zeros((y_valid_c_hot.shape[0], 1))

        if self.args.training_type == 'domain_adaptation':
            print('Applying Domain Adaptation')
            # Analysing the target train dataset
            central_pixels_coor_tr = self.dataset_t.central_pixels_coor_tr.copy()
            y_train = self.dataset_t.y_train.copy()
            central_pixels_coor_tr, y_train = shuffle(central_pixels_coor_tr, y_train, random_state=0)
            positive_coordinates = np.transpose(np.array(np.where(y_train == 1)))
            negative_coordinates = np.transpose(np.array(np.where(y_train == 0)))
            positive_coordinates = positive_coordinates[:,0]
            negative_coordinates = negative_coordinates[:,0]
            if len(negative_coordinates) != 0:
                positive_central_pixels_coor_tr = central_pixels_coor_tr[positive_coordinates, :]
                if self.args.data_augmentation:
                    positive_central_pixels_coor_tr, _ = Data_Augmentation_Definition(positive_central_pixels_coor_tr, np.ones((len(positive_coordinates),1)))

                #Taking the same amount of negative samples as positive
                negative_coordinates = negative_coordinates[:positive_central_pixels_coor_tr.shape[0]]
                if self.args.data_augmentation:
                    negative_central_pixels_coor_tr = np.concatenate((central_pixels_coor_tr[negative_coordinates, :], np.zeros((len(negative_coordinates),1))),axis=1)
                else:
                    negative_central_pixels_coor_tr = central_pixels_coor_tr[negative_coordinates, :]

                positive_y_train = np.ones((positive_central_pixels_coor_tr.shape[0],1))
                negative_y_train = np.zeros((negative_central_pixels_coor_tr.shape[0],1))
                central_pixels_coor_tr = np.concatenate((positive_central_pixels_coor_tr, negative_central_pixels_coor_tr), axis=0)
                y_train = np.concatenate((positive_y_train, negative_y_train), axis=0)
                # Shuffling again
                central_pixels_coor_tr_t, y_train_t = shuffle(central_pixels_coor_tr, y_train,random_state=0)
            else:
                positive_central_pixels_coor_tr = central_pixels_coor_tr[positive_coordinates, :]
                if self.args.data_augmentation:
                    positive_central_pixels_coor_tr, _ = Data_Augmentation_Definition(positive_central_pixels_coor_tr, np.ones((len(positive_coordinates),1)))

                central_pixels_coor_tr_t = positive_central_pixels_coor_tr.copy()
                y_train_t = np.ones((positive_central_pixels_coor_tr.shape[0],1))

            # Analysing the target valid dataset
            central_pixels_coor_vl = self.dataset_t.central_pixels_coor_vl.copy()
            y_valid = self.dataset_t.y_valid.copy()
            central_pixels_coor_vl, y_valid = shuffle(central_pixels_coor_vl, y_valid, random_state=0)
            positive_coordinates = np.transpose(np.array(np.where(y_valid == 1)))
            negative_coordinates = np.transpose(np.array(np.where(y_valid == 0)))
            positive_coordinates = positive_coordinates[:,0]
            negative_coordinates = negative_coordinates[:,0]
            if len(negative_coordinates) != 0:
                positive_central_pixels_coor_vl = central_pixels_coor_vl[positive_coordinates, :]
                if self.args.data_augmentation:
                    positive_central_pixels_coor_vl, _ = Data_Augmentation_Definition(positive_central_pixels_coor_vl, np.ones((len(positive_coordinates),1)))

                #Taking the same amount of negative samples as positive
                negative_coordinates = negative_coordinates[:positive_central_pixels_coor_vl.shape[0]]
                if self.args.data_augmentation:
                    negative_central_pixels_coor_vl = np.concatenate((central_pixels_coor_vl[negative_coordinates, :], np.zeros((len(negative_coordinates),1))),axis=1)
                else:
                    negative_central_pixels_coor_vl = central_pixels_coor_vl[negative_coordinates, :]

                positive_y_valid = np.ones((positive_central_pixels_coor_vl.shape[0],1))
                negative_y_valid = np.zeros((negative_central_pixels_coor_vl.shape[0],1))
                central_pixels_coor_vl = np.concatenate((positive_central_pixels_coor_vl, negative_central_pixels_coor_vl), axis=0)
                y_valid = np.concatenate((positive_y_valid, negative_y_valid), axis=0)
                # Shuffling again
                central_pixels_coor_vl_t, y_valid_t = shuffle(central_pixels_coor_vl, y_valid,random_state=0)
            else:
                positive_central_pixels_coor_vl = central_pixels_coor_vl[positive_coordinates, :]
                if self.args.data_augmentation:
                    positive_central_pixels_coor_vl, _ = Data_Augmentation_Definition(positive_central_pixels_coor_vl, np.ones((len(positive_coordinates),1)))

                central_pixels_coor_vl_t = positive_central_pixels_coor_vl.copy()
                y_valid_t = np.ones((positive_central_pixels_coor_vl.shape[0],1))
            print("Target sets dimensions")
            print(np.shape(central_pixels_coor_tr_t))
            print(np.shape(central_pixels_coor_vl_t))
            print(np.shape(y_train_t))
            print(np.shape(y_valid_t))

            #Verify the size of each set aiming at balancing both training sets
            size_s = np.size(y_train_s, 0)
            size_t = np.size(y_train_t, 0)
            if size_t > size_s:
                positive_coordinates = np.transpose(np.array(np.where(y_train_t == 1)))
                negative_coordinates = np.transpose(np.array(np.where(y_train_t == 0)))
                positive_coordinates = positive_coordinates[:,0]
                negative_coordinates = negative_coordinates[:,0]
                if len(negative_coordinates) != 0:
                    central_pixels_coor_tr_p = central_pixels_coor_tr_t[positive_coordinates,:]
                    central_pixels_coor_tr_n = central_pixels_coor_tr_t[negative_coordinates,:]
                    y_train_p = y_train_t[positive_coordinates,:]
                    y_train_n = y_train_t[negative_coordinates,:]
                    central_pixels_coor_tr_p = central_pixels_coor_tr_p[:int(size_s/2),:]
                    central_pixels_coor_tr_n = central_pixels_coor_tr_n[:int(size_s/2),:]
                    y_train_p = y_train_p[:int(size_s/2),:]
                    y_train_n = y_train_n[:int(size_s/2),:]
                    central_pixels_coor_tr = np.concatenate((central_pixels_coor_tr_p, central_pixels_coor_tr_n), axis=0)
                    y_train = np.concatenate((y_train_p, y_train_n), axis=0)
                    central_pixels_coor_tr_t, y_train_t = shuffle(central_pixels_coor_tr, y_train,random_state=0)
                else:
                    central_pixels_coor_tr_t = central_pixels_coor_tr_t[:size_s, :]
                    y_train_t = y_train_t[:size_s, :]
            elif size_s > size_t:
                positive_coordinates = np.transpose(np.array(np.where(y_train_s == 1)))
                negative_coordinates = np.transpose(np.array(np.where(y_train_s == 0)))
                positive_coordinates = positive_coordinates[:,0]
                negative_coordinates = negative_coordinates[:,0]

                central_pixels_coor_tr_p = central_pixels_coor_tr_s[positive_coordinates,:]
                central_pixels_coor_tr_n = central_pixels_coor_tr_s[negative_coordinates,:]
                y_train_p = y_train_s[positive_coordinates,:]
                y_train_n = y_train_s[negative_coordinates,:]
                central_pixels_coor_tr_p = central_pixels_coor_tr_p[:int(size_t/2),:]
                central_pixels_coor_tr_n = central_pixels_coor_tr_n[:int(size_t/2),:]
                y_train_p = y_train_p[:int(size_t/2),:]
                y_train_n = y_train_n[:int(size_t/2),:]
                central_pixels_coor_tr = np.concatenate((central_pixels_coor_tr_p, central_pixels_coor_tr_n), axis=0)
                y_train = np.concatenate((y_train_p, y_train_n), axis=0)
                central_pixels_coor_tr_s, y_train_s = shuffle(central_pixels_coor_tr, y_train,random_state=0)

            #Verify the size of each set aiming at balancing both validation sets
            size_s = np.size(y_valid_s, 0)
            size_t = np.size(y_valid_t, 0)
            if size_t > size_s:
                positive_coordinates = np.transpose(np.array(np.where(y_valid_t == 1)))
                negative_coordinates = np.transpose(np.array(np.where(y_valid_t == 0)))
                positive_coordinates = positive_coordinates[:,0]
                negative_coordinates = negative_coordinates[:,0]
                if len(negative_coordinates) != 0:
                    central_pixels_coor_vl_p = central_pixels_coor_vl_t[positive_coordinates,:]
                    central_pixels_coor_vl_n = central_pixels_coor_vl_t[negative_coordinates,:]
                    y_valid_p = y_valid_t[positive_coordinates,:]
                    y_valid_n = y_valid_t[negative_coordinates,:]
                    central_pixels_coor_vl_p = central_pixels_coor_vl_p[:int(size_s/2),:]
                    central_pixels_coor_vl_n = central_pixels_coor_vl_n[:int(size_s/2),:]
                    y_valid_p = y_valid_p[:int(size_s/2),:]
                    y_valid_n = y_valid_n[:int(size_s/2),:]
                    central_pixels_coor_vl = np.concatenate((central_pixels_coor_vl_p, central_pixels_coor_vl_n), axis=0)
                    y_valid = np.concatenate((y_valid_p, y_valid_n), axis=0)
                    central_pixels_coor_vl_t, y_valid_t = shuffle(central_pixels_coor_vl, y_valid, random_state=0)
                else:
                    central_pixels_coor_vl_t = central_pixels_coor_vl_t[:size_s, :]
                    y_valid_t = y_valid_t[:size_s, :]

            elif size_s > size_t:
                positive_coordinates = np.transpose(np.array(np.where(y_valid_s == 1)))
                negative_coordinates = np.transpose(np.array(np.where(y_valid_s == 0)))
                positive_coordinates = positive_coordinates[:,0]
                negative_coordinates = negative_coordinates[:,0]

                central_pixels_coor_vl_p = central_pixels_coor_vl_s[positive_coordinates,:]
                central_pixels_coor_vl_n = central_pixels_coor_vl_s[negative_coordinates,:]
                y_valid_p = y_valid_s[positive_coordinates,:]
                y_valid_n = y_valid_s[negative_coordinates,:]
                central_pixels_coor_vl_p = central_pixels_coor_vl_p[:int(size_t/2),:]
                central_pixels_coor_vl_n = central_pixels_coor_vl_n[:int(size_t/2),:]
                y_valid_p = y_valid_p[:int(size_t/2),:]
                y_valid_n = y_valid_n[:int(size_t/2),:]
                central_pixels_coor_vl = np.concatenate((central_pixels_coor_vl_p, central_pixels_coor_vl_n), axis=0)
                y_valid = np.concatenate((y_valid_p, y_valid_n), axis=0)
                central_pixels_coor_vl_s, y_valid_s = shuffle(central_pixels_coor_vl, y_valid,random_state=0)

            print("Source and Target dimensions")
            print(np.shape(central_pixels_coor_tr_s))
            print(np.shape(central_pixels_coor_tr_t))
            print(np.shape(central_pixels_coor_vl_s))
            print(np.shape(central_pixels_coor_vl_t))
            print(np.shape(y_train_s))
            print(np.shape(y_train_t))
            print(np.shape(y_valid_s))
            print(np.shape(y_valid_t))

            #Preparing the sets for the training
            y_train_ds = np.zeros((y_train_s.shape[0], 1))
            y_valid_ds = np.zeros((y_valid_s.shape[0], 1))
            y_train_dt = np.ones((y_train_t.shape[0], 1))
            y_valid_dt = np.ones((y_valid_t.shape[0], 1))

            y_train_s_hot = tf.keras.utils.to_categorical(y_train_s, self.args.num_classes)
            y_valid_s_hot = tf.keras.utils.to_categorical(y_valid_s, self.args.num_classes)
            y_train_t_hot = tf.keras.utils.to_categorical(y_train_t, self.args.num_classes)
            y_valid_t_hot = tf.keras.utils.to_categorical(y_valid_t, self.args.num_classes)
            y_train_ds_hot = tf.keras.utils.to_categorical(y_train_ds, self.args.num_classes)
            y_valid_ds_hot = tf.keras.utils.to_categorical(y_valid_ds, self.args.num_classes)
            y_train_dt_hot = tf.keras.utils.to_categorical(y_train_dt, self.args.num_classes)
            y_valid_dt_hot = tf.keras.utils.to_categorical(y_valid_dt, self.args.num_classes)

            central_pixels_coor_tr = np.concatenate((central_pixels_coor_tr_s, central_pixels_coor_tr_t), axis = 0)
            central_pixels_coor_vl = np.concatenate((central_pixels_coor_vl_s, central_pixels_coor_vl_t), axis = 0)
            y_train_c_hot = np.concatenate((y_train_s_hot, y_train_t_hot), axis = 0)
            y_valid_c_hot = np.concatenate((y_valid_s_hot, y_valid_t_hot), axis = 0)
            y_train_d_hot = np.concatenate((y_train_ds_hot, y_train_dt_hot), axis = 0)
            y_valid_d_hot = np.concatenate((y_valid_ds_hot, y_valid_dt_hot), axis = 0)
            domain_indexs_tr = np.concatenate((y_train_ds, y_train_dt), axis = 0)
            domain_indexs_vl = np.concatenate((y_valid_ds, y_valid_dt), axis = 0)

            classification_mask_tr = np.concatenate((np.ones((y_train_ds.shape[0] , 1)), np.zeros((y_train_dt.shape[0] , 1))),axis = 0)
            classification_mask_vl = np.concatenate((np.ones((y_valid_ds.shape[0] , 1)), np.zeros((y_valid_dt.shape[0] , 1))),axis = 0)
            if 'CL' in self.args.da_type:
                classification_mask_tr = np.ones((domain_indexs_tr.shape[0] , 1))
                classification_mask_vl = np.ones((domain_indexs_vl.shape[0] , 1))

        data = []
        x_train_s = np.concatenate((self.dataset_s.images_norm[0], self.dataset_s.images_norm[1]), axis = 2)
        data.append(x_train_s)
        if self.args.training_type == 'domain_adaptation':
            x_train_t = np.concatenate((self.dataset_t.images_norm[0], self.dataset_t.images_norm[1]), axis = 2)
            data.append(x_train_t)

        #Computing the number of batches
        num_batches_tr = central_pixels_coor_tr.shape[0]//self.args.batch_size
        num_batches_vl = central_pixels_coor_vl.shape[0]//self.args.batch_size

        e = 0
        while (e < self.args.epochs):

            #Shuffling the data and the labels
            print(central_pixels_coor_tr.shape[0])
            print(y_train_c_hot.shape[0])
            num_samples = central_pixels_coor_tr.shape[0]
            index = np.arange(num_samples)
            np.random.shuffle(index)
            central_pixels_coor_tr = central_pixels_coor_tr[index, :]
            y_train_c_hot = y_train_c_hot[index, :]
            y_train_d_hot = y_train_d_hot[index, :]
            classification_mask_tr = classification_mask_tr[index, :]

            domain_indexs_tr = domain_indexs_tr[index, :]

            num_samples = central_pixels_coor_vl.shape[0]
            index = np.arange(num_samples)
            np.random.shuffle(index)
            central_pixels_coor_vl = central_pixels_coor_vl[index, :]
            y_valid_c_hot = y_valid_c_hot[index, :]
            y_valid_d_hot = y_valid_d_hot[index, :]
            classification_mask_vl = classification_mask_vl[index, :]

            domain_indexs_vl = domain_indexs_vl[index, :]

            # Open a file in order to save the training history
            f = open(self.args.save_checkpoint_path + "Log.txt","a")
            #Initializing loss metrics
            loss_cl_tr = np.zeros((1 , 2))
            loss_cl_vl = np.zeros((1 , 2))
            loss_dr_tr = np.zeros((1 , 2))
            loss_dr_vl = np.zeros((1 , 2))

            accuracy_tr = 0
            f1_score_tr = 0
            recall_tr = 0
            precission_tr = 0

            accuracy_vl = 0
            f1_score_vl = 0
            recall_vl = 0
            precission_vl = 0

            #Computing some parameters
            self.p = float(e) / self.args.epochs
            self.l = 2. / (1. + np.exp(-10. * self.p)) - 1
            self.lr = self.Learning_rate_decay()
            print(self.p)
            print(self.lr)
            print(self.l)

            batch_counter_cl = 0
            batchs = trange(num_batches_tr)

            for b in batchs:
                central_pixels_coor_tr_batch = central_pixels_coor_tr[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                domain_index_batch = domain_indexs_tr[b * self.args.batch_size : (b + 1) * self.args.batch_size, :]
                classification_mask_batch = classification_mask_tr[b * self.args.batch_size : (b + 1) * self.args.batch_size, :]

                if self.args.data_augmentation:
                    transformation_indexs_batch = central_pixels_coor_tr[b * self.args.batch_size : (b + 1) * self.args.batch_size , 2]

                y_train_c_hot_batch = y_train_c_hot[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                y_train_d_hot_batch = y_train_d_hot[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                #Extracting the data patches from it's coordinates
                data_batch = Patch_Extraction(data, central_pixels_coor_tr_batch, domain_index_batch, self.args.patches_dimension, True, 'reflect')
                # Perform data augmentation?
                if self.args.data_augmentation:
                    data_batch = Data_Augmentation_Execution(data_batch, transformation_indexs_batch)

                if self.args.training_type == 'classification':
                    _, batch_loss, batch_probs = self.sess.run([self.training_optimizer, self.classifier_loss, self.prediction_c],
                                                                    feed_dict={self.data: data_batch, self.label: y_train_c_hot_batch,
                                                                    self.mask_c: classification_mask_batch[:,0], self.learning_rate: self.lr})
                    loss_cl_tr[0 , 0] += batch_loss
                    y_train_predict_batch = np.argmax(batch_probs, axis = 1)
                    y_train_batch = np.argmax(y_train_c_hot_batch, axis = 1)
                    batch_counter_cl += 1

                    accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_train_batch.astype(int), y_train_predict_batch.astype(int))

                    accuracy_tr += accuracy
                    f1_score_tr += f1score
                    recall_tr += recall
                    precission_tr += precission
                if self.args.training_type == 'domain_adaptation':
                    if np.sum(domain_index_batch) < self.args.batch_size:
                        _, batch_loss, batch_probs, batch_loss_d = self.sess.run([self.training_optimizer, self.classifier_loss, self.prediction_c, self.domainregressor_loss],
                                                                                feed_dict = {self.data: data_batch, self.label: y_train_c_hot_batch, self.label_d: y_train_d_hot_batch,
                                                                                self.mask_c: classification_mask_batch[:,0], self.L: self.l, self.learning_rate: self.lr})
                        loss_dr_tr[0 , 0] += batch_loss_d
                        loss_cl_tr[0 , 0] += batch_loss

                        #Choosing the source samples only
                        labeled_coordinates = np.transpose(np.array(np.where(domain_index_batch == 0)))
                        labeled_coordinates = labeled_coordinates[:,0]
                        y_train_hot_labeled = y_train_c_hot_batch[labeled_coordinates, :]
                        p_train_labeled_data = batch_probs[labeled_coordinates, :]

                        y_train_predict_batch = np.argmax(p_train_labeled_data, axis = 1)
                        y_train_batch = np.argmax(y_train_hot_labeled, axis = 1)
                        batch_counter_cl += 1

                        accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_train_batch.astype(int), y_train_predict_batch.astype(int))

                        accuracy_tr += accuracy
                        f1_score_tr += f1score
                        recall_tr += recall
                        precission_tr += precission

            loss_cl_tr = loss_cl_tr/batch_counter_cl
            accuracy_tr = accuracy_tr/batch_counter_cl
            f1_score_tr = f1_score_tr/batch_counter_cl
            recall_tr = recall_tr/batch_counter_cl
            precission_tr = precission_tr/batch_counter_cl
            print(batch_counter_cl)
            if self.args.training_type == 'domain_adaptation' and 'DR' in self.args.da_type:
                loss_dr_tr = loss_dr_tr/batch_counter_cl
                print ("%d [Tr loss: %f, acc.: %.2f%%,  precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%, Dr loss: %f]" % (e, loss_cl_tr[0 , 0], accuracy_tr, precission_tr, recall_tr, f1_score_tr, loss_dr_tr[0,0]))
                f.write("%d [Tr loss: %f, acc.: %.2f%%, precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%, Dr loss: %f]\n" % (e, loss_cl_tr[0 , 0], accuracy_tr, precission_tr, recall_tr, f1_score_tr, loss_dr_tr[0,0]))
            else:
                print ("%d [Tr loss: %f, acc.: %.2f%%, precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%]" % (e, loss_cl_tr[0,0], accuracy_tr, precission_tr, recall_tr, f1_score_tr))
                f.write("%d [Tr loss: %f, acc.: %.2f%%, precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%]\n" % (e, loss_cl_tr[0,0], accuracy_tr, precission_tr, recall_tr, f1_score_tr))

            #Computing the validation loss
            print('[*]Computing the validation loss...')
            batch_counter_cl = 0
            batchs = trange(num_batches_vl)

            for b in batchs:
                central_pixels_coor_vl_batch = central_pixels_coor_vl[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                domain_index_batch = domain_indexs_vl[b * self.args.batch_size : (b + 1) * self.args.batch_size, :]
                classification_mask_batch = classification_mask_vl[b * self.args.batch_size : (b + 1) * self.args.batch_size, :]
                if self.args.data_augmentation:
                    transformation_indexs_batch = central_pixels_coor_vl[b * self.args.batch_size : (b + 1) * self.args.batch_size , 2]

                y_valid_c_hot_batch = y_valid_c_hot[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                y_valid_d_hot_batch = y_valid_d_hot[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
                #Extracting the data patches from it's coordinates
                data_batch = Patch_Extraction(data, central_pixels_coor_vl_batch, domain_index_batch, self.args.patches_dimension, True, 'reflect')

                if self.args.data_augmentation:
                    data_batch = Data_Augmentation_Execution(data_batch, transformation_indexs_batch)

                if self.args.training_type == 'classification':
                    batch_loss, batch_probs = self.sess.run([self.classifier_loss, self.prediction_c],
                                                                    feed_dict={self.data: data_batch, self.label: y_valid_c_hot_batch,
                                                                    self.mask_c: classification_mask_batch[:,0], self.learning_rate: self.lr})
                    loss_cl_vl[0 , 0] += batch_loss
                    y_valid_batch = np.argmax(y_valid_c_hot_batch, axis = 1)
                    y_valid_predict_batch = np.argmax(batch_probs, axis = 1)
                    accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_valid_batch.astype(int), y_valid_predict_batch.astype(int))
                    accuracy_vl += accuracy
                    f1_score_vl += f1score
                    recall_vl += recall
                    precission_vl += precission
                    batch_counter_cl += 1

                if self.args.training_type == 'domain_adaptation':
                    if np.sum(domain_index_batch) < self.args.batch_size:
                        batch_loss, batch_probs, batch_loss_d = self.sess.run([self.classifier_loss, self.prediction_c, self.domainregressor_loss],
                                                                                feed_dict = {self.data: data_batch, self.label: y_valid_c_hot_batch, self.label_d: y_valid_d_hot_batch,
                                                                                self.mask_c: classification_mask_batch[:,0], self.L: self.l, self.learning_rate: self.lr})
                        loss_dr_vl[0 , 0] += batch_loss_d
                        loss_cl_vl[0 , 0] += batch_loss
                        #Choosing the source samples only
                        labeled_coordinates = np.transpose(np.array(np.where(domain_index_batch == 0)))
                        labeled_coordinates = labeled_coordinates[:,0]
                        y_valid_hot_labeled = y_valid_c_hot_batch[labeled_coordinates, :]
                        p_valid_labeled_data = batch_probs[labeled_coordinates, :]

                        y_valid_batch = np.argmax(y_valid_hot_labeled, axis = 1)
                        y_valid_predict_batch = np.argmax(p_valid_labeled_data, axis=1)
                        batch_counter_cl += 1

                        accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_valid_batch.astype(int), y_valid_predict_batch.astype(int))

                        accuracy_vl += accuracy
                        f1_score_vl += f1score
                        recall_vl += recall
                        precission_vl += precission

            loss_cl_vl = loss_cl_vl/(batch_counter_cl)
            accuracy_vl = accuracy_vl/(batch_counter_cl)
            f1_score_vl = f1_score_vl/(batch_counter_cl)
            recall_vl = recall_vl/(batch_counter_cl)
            precission_vl = precission_vl/(batch_counter_cl)

            if self.args.training_type == 'domain_adaptation' and 'DR' in self.args.da_type:
                loss_dr_vl = loss_dr_vl/batch_counter_cl
                print ("%d [Vl loss: %f, acc.: %.2f%%,  precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%, DrV loss: %f]" % (e, loss_cl_vl[0,0], accuracy_vl, precission_vl, recall_vl, f1_score_vl, loss_dr_vl[0 , 0]))
                f.write("%d [Vl loss: %f, acc.: %.2f%%, precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%, DrV loss: %f]\n" % (e, loss_cl_vl[0,0], accuracy_vl, precission_vl, recall_vl, f1_score_vl, loss_dr_vl[0 , 0]))

            else:
                print ("%d [Vl loss: %f, acc.: %.2f%%,  precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%]" % (e, loss_cl_vl[0,0], accuracy_vl, precission_vl, recall_vl, f1_score_vl))
                f.write("%d [Vl loss: %f, acc.: %.2f%%, precission: %.2f%%, recall: %.2f%%, fscore: %.2f%%]\n" % (e, loss_cl_vl[0,0], accuracy_vl, precission_vl, recall_vl, f1_score_vl))

            if self.args.training_type == 'domain_adaptation' and 'DR' in self.args.da_type:
                if self.l != 0:
                    FLAG = False
                    if  best_val_dr < loss_dr_vl[0 , 0] and loss_dr_vl[0 , 0] < 1:
                        if best_val_fs < f1_score_vl:
                            best_val_dr = loss_dr_vl[0 , 0]
                            best_val_fs = f1_score_vl
                            best_mod_fs = f1_score_vl
                            best_mod_dr = loss_dr_vl[0 , 0]
                            best_model_epoch = e
                            print('[!]Saving best ideal model at epoch: ' + str(e))
                            f.write("[!]Ideal best ideal model\n")
                            self.save(self.args.save_checkpoint_path, best_model_epoch)
                            FLAG = True
                        elif np.abs(best_val_fs - f1_score_vl) < 3:
                            best_val_dr = loss_dr_vl[0 , 0]
                            best_mod_fs = f1_score_vl
                            best_mod_dr = loss_dr_vl[0 , 0]
                            best_model_epoch = e
                            print('[!]Saving best model attending best Dr_loss at epoch: ' + str(e))
                            f.write("[!]Best model attending best Dr_loss\n")
                            self.save(self.args.save_checkpoint_path, best_model_epoch)
                            FLAG = True
                    elif best_val_fs < f1_score_vl:
                        if  np.abs(best_val_dr - loss_dr_vl[0 , 0]) < 0.2:
                            best_val_fs = f1_score_vl
                            best_mod_fs = f1_score_vl
                            best_mod_dr = loss_dr_vl[0 , 0]
                            best_model_epoch = e
                            print('[!]Saving best model attending best f1-score at epoch: ' + str(e))
                            f.write("[!]Best model attending best f1-score \n")
                            self.save(self.args.save_checkpoint_path, best_model_epoch)
                            FLAG = True
                    if FLAG:
                        pat = 0
                        print('[!] Best Model with DrV loss: %.3f and F1-Score: %.2f%%'% (best_mod_dr, best_mod_fs))
                    else:
                        print('[!] The Model has not been considered as suitable for saving procedure.')
                        pat += 1
                        if pat > self.args.patience:
                            break
                else:
                    print('Warming Up!')
            else:
                if best_val_fs < f1_score_vl:
                    best_val_fs = f1_score_vl
                    pat = 0
                    print('[!]Saving best model ...')
                    f.write('[!]Saving best model...')
                    self.save(self.args.save_checkpoint_path, e)
                else:
                    pat += 1
                    if pat > self.args.patience:
                        break
            f.close()
            e += 1

    def Test(self):

        print(np.shape(self.dataset.images_norm[0]))
        print(np.shape(self.dataset.images_norm[1]))
        heat_map = np.zeros((self.dataset.images_norm[0].shape[0] + 2 * (self.args.patches_dimension//2), self.dataset.images_norm[0].shape[1] + 2 * (self.args.patches_dimension//2)))
        print(np.shape(heat_map))
        x_test = []
        data = np.concatenate((self.dataset.images_norm[0], self.dataset.images_norm[1]), axis = 2)
        x_test.append(data)

        num_batches_ts = self.dataset.central_pixels_coor_ts.shape[0]//self.args.batch_size
        batchs = trange(num_batches_ts)
        print(num_batches_ts)

        for b in batchs:
            self.central_pixels_coor_ts_batch = self.dataset.central_pixels_coor_ts[b * self.args.batch_size : (b + 1) * self.args.batch_size , :]
            self.x_test_batch = Patch_Extraction(x_test, self.central_pixels_coor_ts_batch, np.zeros((self.args.batch_size , 1)), self.args.patches_dimension, True, 'reflect')

            probs = self.sess.run(self.prediction_c,
                                         feed_dict={self.data: self.x_test_batch})

            for i in range(self.args.batch_size):
                heat_map[int(self.central_pixels_coor_ts_batch[i,0]), int(self.central_pixels_coor_ts_batch[i,1])] = probs[i,1]

        if (num_batches_ts * self.args.batch_size) < self.dataset.central_pixels_coor_ts.shape[0]:
            self.central_pixels_coor_ts_batch = self.dataset.central_pixels_coor_ts[num_batches_ts * self.args.batch_size : , :]
            self.x_test_batch = Patch_Extraction(x_test, self.central_pixels_coor_ts_batch, np.zeros((self.central_pixels_coor_ts_batch.shape[0] , 1)), self.args.patches_dimension, True, 'reflect')

            probs = self.sess.run(self.prediction_c,
                                         feed_dict={self.data: self.x_test_batch})
            for i in range(self.central_pixels_coor_ts_batch.shape[0]):
                heat_map[int(self.central_pixels_coor_ts_batch[i,0]), int(self.central_pixels_coor_ts_batch[i,1])] = probs[i,1]

        np.save(self.args.save_results_dir + 'heat_map', heat_map)

    def save(self, checkpoint_dir, epoch):
        model_name = "DANN.model"

        self.saver.save(self.sess,
                        os.path.join(checkpoint_dir, model_name),
                        global_step=epoch)
        print("Checkpoint Saved with SUCCESS!")

    def load(self, checkpoint_dir):
        print(" [*] Reading checkpoint...")
        print(checkpoint_dir)

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            aux = 'model_example'
            for i in range(len(ckpt_name)):
                if ckpt_name[-i-1] == '-':
                    aux = ckpt_name[-i:]
                    break
            return aux
        else:
            return ''

def Metrics_For_Test(heat_map,
                     reference_t1, reference_t2,
                     Train_tiles, Valid_tiles, Undesired_tiles,
                     Thresholds,
                     args):

    half_dim = args.patches_dimension//2

    save_path = args.results_dir + args.file + '/'
    print('[*]Defining the initial central patches coordinates...')
    mask_init = mask_creation(reference_t1.shape[0], reference_t1.shape[1], args.horizontal_blocks, args.vertical_blocks, [], [], [])
    mask_final = mask_creation(reference_t1.shape[0], reference_t1.shape[1], args.horizontal_blocks, args.vertical_blocks, Train_tiles, Valid_tiles, Undesired_tiles)

    #mask_final = mask_final_.copy()
    mask_final[mask_final == 1] = 0
    mask_final[mask_final == 3] = 0
    mask_final[mask_final == 2] = 1
    heat_map = heat_map[half_dim : -half_dim, half_dim : -half_dim]
    Probs_init = heat_map
    positive_map_init = np.zeros_like(Probs_init)

    # Metrics containers
    ACCURACY = np.zeros((1, len(Thresholds)))
    FSCORE = np.zeros((1, len(Thresholds)))
    RECALL = np.zeros((1, len(Thresholds)))
    PRECISSION = np.zeros((1, len(Thresholds)))
    CONFUSION_MATRIX = np.zeros((2 , 2, len(Thresholds)))
    CLASSIFICATION_MAPS = np.zeros((len(Thresholds), heat_map.shape[0], heat_map.shape[1], 3))
    ALERT_AREA = np.zeros((1 , len(Thresholds)))

    print('[*]The metrics computation has started...')
    #Computing the metrics for each defined threshold
    for th in range(len(Thresholds)):
        print(Thresholds[th])

        positive_map_init = np.zeros_like(heat_map)
        reference_t1_copy = reference_t1.copy()

        threshold = Thresholds[th]
        positive_coordinates = np.transpose(np.array(np.where(Probs_init >= threshold)))
        positive_map_init[positive_coordinates[:,0].astype('int'), positive_coordinates[:,1].astype('int')] = 1

        if args.eliminate_regions:
            positive_map_init_ = skimage.morphology.area_opening(positive_map_init.astype('int'),area_threshold = args.area_avoided, connectivity=1)
            eliminated_samples = positive_map_init - positive_map_init_
        else:
            eliminated_samples = np.zeros_like(heat_map)

        reference_t1_copy = reference_t1_copy + eliminated_samples
        reference_t1_copy[reference_t1_copy == 2] = 1

        reference_t1_copy = reference_t1_copy - 1
        reference_t1_copy[reference_t1_copy == -1] = 1
        reference_t1_copy[reference_t2 == 2] = 0
        mask_f = mask_final * reference_t1_copy

        central_pixels_coordinates_ts_ = np.transpose(np.array(np.where(mask_f == 1)))
        y_test = reference_t2[central_pixels_coordinates_ts_[:,0].astype('int'), central_pixels_coordinates_ts_[:,1].astype('int')]

        Probs = heat_map[central_pixels_coordinates_ts_[:,0].astype('int'), central_pixels_coordinates_ts_[:,1].astype('int')]
        Probs[Probs >= Thresholds[th]] = 1
        Probs[Probs <  Thresholds[th]] = 0

        accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_test.astype('int'), Probs.astype('int'))

        #Classification_map, _, _ = Classification_Maps(Probs, y_test, central_pixels_coordinates_ts_, hit_map)

        TP = conf_mat[1 , 1]
        FP = conf_mat[0 , 1]
        TN = conf_mat[0 , 0]
        FN = conf_mat[1 , 0]
        numerator = TP + FP

        denominator = TN + FN + FP + TP

        Alert_area = 100*(numerator/denominator)
        print(f1score)
        ACCURACY[0 , th] = accuracy
        FSCORE[0 , th] = f1score
        RECALL[0 , th] = recall
        PRECISSION[0 , th] = precission
        CONFUSION_MATRIX[: , : , th] = conf_mat
        #CLASSIFICATION_MAPS[th, :, :, :] = Classification_map
        ALERT_AREA[0 , th] = Alert_area

    #Saving the metrics as npy array
    if not args.save_result_text:
        np.save(save_path + 'Accuracy', ACCURACY)
        np.save(save_path + 'Fscore', FSCORE)
        np.save(save_path + 'Recall', RECALL)
        np.save(save_path + 'Precission', PRECISSION)
        np.save(save_path + 'Confusion_matrix', CONFUSION_MATRIX)
        #np.save(save_path + 'Classification_maps', CLASSIFICATION_MAPS)
        np.save(save_path + 'Alert_area', ALERT_AREA)

    print('Accuracy')
    print(ACCURACY)
    print('Fscore')
    print(FSCORE)
    print('Recall')
    print(RECALL)
    print('Precision')
    print(PRECISSION)
    print('Confusion matrix')
    print(CONFUSION_MATRIX[:,:,0])
    print('Alert_area')
    print(ALERT_AREA)

    return ACCURACY, FSCORE, RECALL, PRECISSION, CONFUSION_MATRIX, ALERT_AREA

def Metrics_For_Test_M(heat_map,
                     reference_t1, reference_t2,
                     Train_tiles, Valid_tiles, Undesired_tiles,
                     args):

    half_dim = args.patches_dimension//2

    save_path = args.results_dir + args.file + '/'
    print('[*]Defining the initial central patches coordinates...')
    mask_init = mask_creation(reference_t1.shape[0], reference_t1.shape[1], args.horizontal_blocks, args.vertical_blocks, [], [], [])
    mask_final = mask_creation(reference_t1.shape[0], reference_t1.shape[1], args.horizontal_blocks, args.vertical_blocks, Train_tiles, Valid_tiles, Undesired_tiles)

    #mask_final = mask_final_.copy()
    mask_final[mask_final == 1] = 0
    mask_final[mask_final == 3] = 0
    mask_final[mask_final == 2] = 1
    heat_map = heat_map[half_dim : -half_dim, half_dim : -half_dim]
    sio.savemat(save_path + 'heat_map.mat' , {'heat_map': heat_map})
    Probs_init = heat_map
    positive_map_init = np.zeros_like(Probs_init)

    reference_t1_copy_ = reference_t1.copy()
    reference_t1_copy_ = reference_t1_copy_ - 1
    reference_t1_copy_[reference_t1_copy_ == -1] = 1
    reference_t1_copy_[reference_t2 == 2] = 0
    mask_f_ = mask_final * reference_t1_copy_
    sio.savemat(save_path + 'mask_f_.mat' , {'mask_f_': mask_f_})
    sio.savemat(save_path + 'reference_t2.mat' , {'reference': reference_t2})
    # Raul Implementation
    min_array = np.zeros((1 , ))
    Pmax = np.max(Probs_init[mask_f_ == 1])
    probs_list = np.arange(Pmax, 0, -Pmax/(args.Npoints - 1))
    Thresholds = np.concatenate((probs_list , min_array))

    print('Max probability value:')
    print(Pmax)
    print('Thresholds:')
    print(Thresholds)
    # Metrics containers
    ACCURACY = np.zeros((1, len(Thresholds)))
    FSCORE = np.zeros((1, len(Thresholds)))
    RECALL = np.zeros((1, len(Thresholds)))
    PRECISSION = np.zeros((1, len(Thresholds)))
    CONFUSION_MATRIX = np.zeros((2 , 2, len(Thresholds)))
    CLASSIFICATION_MAPS = np.zeros((len(Thresholds), heat_map.shape[0], heat_map.shape[1], 3))
    ALERT_AREA = np.zeros((1 , len(Thresholds)))

    print('[*]The metrics computation has started...')
    #Computing the metrics for each defined threshold
    for th in range(len(Thresholds)):
        print(Thresholds[th])

        positive_map_init = np.zeros_like(heat_map)
        reference_t1_copy = reference_t1.copy()

        threshold = Thresholds[th]
        positive_coordinates = np.transpose(np.array(np.where(Probs_init >= threshold)))
        positive_map_init[positive_coordinates[:,0].astype('int'), positive_coordinates[:,1].astype('int')] = 1

        if args.eliminate_regions:
            positive_map_init_ = skimage.morphology.area_opening(positive_map_init.astype('int'),area_threshold = args.area_avoided, connectivity=1)
            eliminated_samples = positive_map_init - positive_map_init_
        else:
            eliminated_samples = np.zeros_like(heat_map)

        reference_t1_copy = reference_t1_copy + eliminated_samples
        reference_t1_copy[reference_t1_copy == 2] = 1
        reference_t1_copy = reference_t1_copy - 1
        reference_t1_copy[reference_t1_copy == -1] = 1
        reference_t1_copy[reference_t2 == 2] = 0
        mask_f = mask_final * reference_t1_copy

        central_pixels_coordinates_ts_ = np.transpose(np.array(np.where(mask_f == 1)))

        y_test = reference_t2[central_pixels_coordinates_ts_[:,0].astype('int'), central_pixels_coordinates_ts_[:,1].astype('int')]

        Probs = heat_map[central_pixels_coordinates_ts_[:,0].astype('int'), central_pixels_coordinates_ts_[:,1].astype('int')]
        Probs[Probs >= Thresholds[th]] = 1
        Probs[Probs <  Thresholds[th]] = 0

        accuracy, f1score, recall, precission, conf_mat = compute_metrics(y_test.astype('int'), Probs.astype('int'))

        #Classification_map, _, _ = Classification_Maps(Probs, y_test, central_pixels_coordinates_ts_, hit_map)

        TP = conf_mat[1 , 1]
        FP = conf_mat[0 , 1]
        TN = conf_mat[0 , 0]
        FN = conf_mat[1 , 0]
        numerator = TP + FP

        denominator = TN + FN + FP + TP

        Alert_area = 100*(numerator/denominator)
        #print(f1score)
        print(precission)
        print(recall)
        ACCURACY[0 , th] = accuracy
        FSCORE[0 , th] = f1score
        RECALL[0 , th] = recall
        PRECISSION[0 , th] = precission
        CONFUSION_MATRIX[: , : , th] = conf_mat
        #CLASSIFICATION_MAPS[th, :, :, :] = Classification_map
        ALERT_AREA[0 , th] = Alert_area

    #Saving the metrics as npy array
    if not args.save_result_text:
        np.save(save_path + 'Accuracy', ACCURACY)
        np.save(save_path + 'Fscore', FSCORE)
        np.save(save_path + 'Recall', RECALL)
        np.save(save_path + 'Precission', PRECISSION)
        np.save(save_path + 'Confusion_matrix', CONFUSION_MATRIX)
        #np.save(save_path + 'Classification_maps', CLASSIFICATION_MAPS)
        np.save(save_path + 'Alert_area', ALERT_AREA)

    print('Accuracy')
    print(ACCURACY)
    print('Fscore')
    print(FSCORE)
    print('Recall')
    print(RECALL)
    print('Precision')
    print(PRECISSION)
    print('Confusion matrix')
    print(CONFUSION_MATRIX[:,:,0])
    print('Alert_area')
    print(ALERT_AREA)

    return ACCURACY, FSCORE, RECALL, PRECISSION, CONFUSION_MATRIX, ALERT_AREA


"""Use PyMOl to create templates for bomeba"""

import numpy as np
np.set_printoptions(precision=2)
import __main__
__main__.pymol_argv = ['pymol', '-qck']
import pymol
from pymol import cmd, stored
pymol.finish_launching()
import openbabel as ob

# set hydrogen names to PDB compliant
cmd.set('pdb_reformat_names_mode', 2)
#cmd.set('retain_order', 1)
sel = 'all'

aa = ['A', 'R', 'N', 'D', 'C', 'E', 'Q', 'G', 'H', 'I', 'L', 'K', 'M',
      'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'B', 'Z']

def minimize(selection='all', forcefield='MMFF94s', method='cg',
             nsteps=2000, conv=1E-8, cutoff=False, cut_vdw=6.0, cut_elec=8.0):
    pdb_string = cmd.get_pdbstr(selection)
    name = cmd.get_legal_name(selection)
    obconversion = ob.OBConversion()
    obconversion.SetInAndOutFormats('pdb', 'pdb')
    mol = ob.OBMol()
    obconversion.ReadString(mol, pdb_string)
    ff = ob.OBForceField.FindForceField(forcefield)
    ff.Setup(mol)
    if cutoff == True:
        ff.EnableCutOff(True)
        ff.SetVDWCutOff(cut_vdw)
        ff.SetElectrostaticCutOff(cut_elec)
    if method == 'cg':
        ff.ConjugateGradients(nsteps, conv)
    else:
        ff.SteepestDescent(nsteps, conv)
    ff.GetCoordinates(mol)
    nrg = ff.Energy()
    pdb_string = obconversion.WriteString(mol)
    cmd.delete(name)
    if name == 'all':
        name = 'all_'
    cmd.read_pdbstr(pdb_string, name)
    return nrg

for res_name in aa:
    # Get coordinates and offset
    # cmd.load('templates/protein/{}.pdb'.format(res_name))
    #stored.IDs = []
    # cmd.iterate('all','stored.IDs.append((ID))')
    #cmd.alter('all', 'ID = ID - stored.IDs[0] - 1')

    if res_name == 'Z':
        cmd.fab('AZ')
        cmd.select('lala', 'resname ala')
        cmd.remove('lala')
    else:
        cmd.fab(res_name)
    nrg = minimize(selection=sel, forcefield='GAFF', method='cg', nsteps=2000)
    # print(nrg)
    xyz = cmd.get_coords(sel)
    offset = len(xyz)

    # get atom names
    stored.atom_names = []
    cmd.iterate(sel, 'stored.atom_names.append(name)')

    # get bonds
    stored.bonds = []
    model = cmd.get_model(sel)
    for at in model.atom:
        cmd.iterate('neighbor ID %s' % at.id,
                    'stored.bonds.append((%s-1, ID-1))' % at.id)
    bonds = list(set([tuple(sorted(i)) for i in stored.bonds]))
    bonds.sort()

    # get bb and sc
    stored.bb = []
    stored.sc = []
    bb = '(nbr (name n+ca+c+o) ) and hydro or (name n+ca+c+o)'
    sc = 'not ({})'.format(bb)
    cmd.iterate(bb, 'stored.bb.append(ID-1)')
    cmd.iterate(sc, 'stored.sc.append(ID-1)')

    # small check before returning the results
    if len(stored.atom_names) == offset:
        if res_name == 'G':
            stored.atom_names = ['N', 'CA', 'C', 'O', 'H', 'HA2', 'HA3']
        res = """{}_info = AA_info(coords=np.{},
             atom_names = {},
             bb = {},
             sc = {},
             bonds = {},
             offset = {})\n""".format(res_name, repr(xyz), stored.atom_names, stored.bb, stored.sc, bonds, offset)
        print(res)
    else:
        print('Something funny is going on here!')
    cmd.delete('all')


# The code is derived from the following repository:
# https://github.com/yeephycho/tensorflow-face-detection

from InputEstimators.FaceDetectors.FaceDetectorABC import FaceDetectorABC
from InputEstimators.FaceDetectors.FaceBox import FaceBox
from paths import TFMobileNetSSDFaceDetector_tf_model_path
import tensorflow as tf
import numpy as np
import cv2

class TFMobileNetSSDFaceDetector(FaceDetectorABC):
    @staticmethod
    def __convert_label_map_to_categories(label_map, max_num_classes, use_display_name=True):
      categories = []
      list_of_ids_already_added = []
      if not label_map:
        label_id_offset = 1
        for class_id in range(max_num_classes):
          categories.append({'id': class_id + label_id_offset, 'name': 'category_{}'.format(class_id + label_id_offset) })
        return categories
      for item in label_map.item:
        if not 0 < item.id <= max_num_classes:
          logging.info('Ignore item %d since it falls outside of requested label range.', item.id)
          continue
        if use_display_name and item.HasField('display_name'):
          name = item.display_name
        else:
          name = item.name
        if item.id not in list_of_ids_already_added:
          list_of_ids_already_added.append(item.id)
          categories.append({'id': item.id, 'name': name})
      return categories

    @staticmethod
    def loadTFGraph(tf_model_path):
        detection_graph = tf.Graph()
        with detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(tf_model_path, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
        return detection_graph

    @staticmethod
    def _decodeFaceBox(detection):
        (rows, cols, _), detection = detection
        x_left_bottom = int(detection[1] * cols)
        y_left_bottom = int(detection[2] * rows)
        x_right_top = int(detection[3] * cols)
        y_right_top = int(detection[0] * rows)
        return FaceBox(x_left_bottom, y_right_top, x_right_top, y_left_bottom)

    def __init__(self, confidence_threshold = 0.90, model_path = None, *args, **kwargs):
        if model_path == None:
            model_path = TFMobileNetSSDFaceDetector_tf_model_path

        self.__graph = self.loadTFGraph(model_path)
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.__sess = tf.Session(graph = self.__graph, config=config)

        super().__init__(*args, **kwargs)

    def _detectFirstFaceBox(self, frame):
        image_tensor = self.__graph.get_tensor_by_name('image_tensor:0')
        boxes = self.__graph.get_tensor_by_name('detection_boxes:0')
        boxes = self.__sess.run([boxes], feed_dict={image_tensor: frame})
        return [b for b in boxes[0][0]]

    def _detectAllFaceBoxesWithScores(self, frame):
        image_tensor = self.__graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        boxes = self.__graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        scores = self.__graph.get_tensor_by_name('detection_scores:0')
        classes = self.__graph.get_tensor_by_name('detection_classes:0')
        num_detections = self.__graph.get_tensor_by_name('num_detections:0')
        outputTensors = [boxes, scores, classes, num_detections]
        # Actual detection.
        outputs = self.__sess.run(outputTensors, feed_dict={image_tensor: frame})
        (boxes, scores, classes, num_detections) = outputs
        return [b for b in boxes[0]]

    def _detectFaceBox(self, frame):
        frameShape = frame.shape
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = np.expand_dims(frame, axis=0)
        #faceBoxDetections = self._detectAllFaceBoxesWithScores(frame)
        faceBoxDetections = self._detectFirstFaceBox(frame)
        faceBoxDetections = [(frameShape, b) for b in faceBoxDetections]
        if len(faceBoxDetections) > 0:
            self._faceBox = self._decodeFaceBox(faceBoxDetections[0])
        return self._faceBox


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

"""This module provides lattice picture from FLAME lattice file(.lat) or FLAME Machine object."""

from flame import Machine
from flame_utils.core import ModelFlame
from flame_utils.core import BeamState
from flame_utils.io import collect_data

import logging
import matplotlib.lines as lin
import matplotlib.patches as ptc
import matplotlib.pyplot as plt
import numpy as np

_LOGGER = logging.getLogger(__name__)

def hplot(*args, **kws):
    """Plot beam state history with input arguments.

    Parameters
    ----------
    machine : object
        FLAME machine or ModelFlame object.
    result :
        List of beam state, generated by Machine.propagete() or ModelFlame.run() method.
    args : str
        Beam state attributes to plot, seprated by comma. (see Note)
    colors : dict
        Color definition of each ``args``.
    lattice : bool
        Plot lattice layout.
    legend : bool
        Plot legend of ``args``
    cs : int
        Index of the charge state to plot.

    Notes
    -----
    List of supported :py:func:`BeamState <flame_utils.core.state.BeamState>` arguments

    pos : float
        Longitudinally propagating position, [m].
    ref_beta : float
        Speed in the unit of light velocity in vacuum of reference charge
        state, Lorentz beta.
    ref_bg : float
        Multiplication of beta and gamma of reference charge state.
    ref_gamma : float
        Relativistic energy of reference charge state, Lorentz gamma.
    ref_IonEk : float
        Kinetic energy of reference charge state, [eV/u].
    ref_IonEs : float
        Rest energy of reference charge state, [eV/u].
    ref_IonQ : int
        Macro particle number of reference charge state.
    ref_IonW : float
        Total energy of reference charge state, [eV/u],
        i.e. :math:`W = E_s + E_k`.
    ref_IonZ : float
        Reference charge to mass ratio, e.g.
        :math:`^{33^{+}}_{238}U: Q[33]/A[238]`.
    ref_phis : float
        Absolute synchrotron phase of reference charge state, [rad].
    ref_SampleIonK : float
        Wave-vector in cavities with different beta values of reference charge
        state.
    beta : Array
        Speed in the unit of light velocity in vacuum of all charge states,
        Lorentz beta.
    bg : Array
        Multiplication of beta and gamma of all charge states.
    gamma : Array
        Relativistic energy of all charge states, Lorentz gamma.
    IonEk : Array
        Kinetic energy of all charge states, [eV/u].
    IonEs : Array
        Rest energy of all charge states, [eV/u].
    IonQ : Array
        Macro particle number of all charge states.
    IonW : Array
        Total energy of all charge states, [eV/u], i.e. :math:`W = E_s + E_k`.
    IonZ : Array
        All charge to mass ratios
    phis : Array
        Absolute synchrotron phase of all charge states, [rad]
    SampleIonK : Array
        Wave-vector in cavities with different beta values of all charge
        states.
    x0_env, xcen : float
        Weight average of all charge states for x', [rad].
    y0_env, ycen : float
        Weight average of all charge states for y, [mm].
    xp0_env, xpcen : float
        Weight average of all charge states for x', [rad].
    yp0_env, ypcen : float
        Weight average of all charge states for y', [rad].
    phi0_env, phicen, zcen: float
        Weight average of all charge states for :math:`\phi`, [rad].
    dEk0_env, dEkcen, zpcen : float
        Weight average of all charge states for :math:`\delta E_k`, [MeV/u].
    x0, xcen_all : Array
        X centroid for all charge states, [mm].
    y0, ycen_all : Array
        Y centroid for all charge states, [mm].
    xp0, xpcen_all : Array
        X centroid divergence for all charge states, [rad].
    yp0, ypcen_all : Array
        Y centroid divergence for all charge states, [rad].
    phi0, phicen_all, zcen_all : Array
        Longitudinal beam length, measured in RF frequency for all charge
        states, [rad].
    dEk0, dEkcen_all, zpcen_all : Array
        Kinetic energy deviation w.r.t. reference charge state, for all charge
        states, [MeV/u].
    x0_rms, xrms : float
        General rms beam envelope for x, [mm].
    y0_rms, yrms : float
        General rms beam envelope for y, [mm].
    xp0_rms, xprms : float
        General rms beam envelope for x', [rad].
    yp0_rms, yprms : float
        General rms beam envelope for y', [rad].
    phi0_rms, phirms, zrms : float
        General rms beam envelope for :math:`\phi`, [rad].
    dEk0_rms, dEkrms, zprms : float
        General rms beam envelope for :math:`\delta E_k`, [MeV/u].
    xrms_all : Array
        General rms beam envelope for x of all charge states, [mm].
    yrms_all : Array
        General rms beam envelope for y of all charge states, [mm].
    xprms_all : Array
        General rms beam envelope for x' of all charge states, [rad].
    yprms_all : Array
        General rms beam envelope for y' of all charge states, [rad].
    phirms_all : Array
        General rms beam envelope for :math:`\phi` of all charge states, [rad].
    dEkrms_all : Array
        General rms beam envelope for :math:`\delta E_k` of all charge states, [MeV/u].
    moment0_env, cenvector : Array
        Weight average of centroid for all charge states, array of
        ``[x, x', y, y', phi, dEk, 1]``, with the units of
        ``[mm, rad, mm, rad, rad, MeV/u, 1]``.
    moment0, cenvector_all : Array
        Centroid for all charge states, array of ``[x, x', y, y', phi, dEk, 1]``.
    moment0_rms, rmsvector : Array
        RMS beam envelope, part of statistical results from ``moment1``.
    moment1, beammatrix_all : Array
        Covariance matrices of all charge states, for each charge state.
    moment1_env, beammatrix : Array
        Covariance matrices of all charge states, average over all charge states.
    xemittance, xeps : float
        Weight average of geometrical x emittance, [mm-mrad].
    yemittance, yeps : float
        Weight average of geometrical y emittance, [mm-mrad].
    zemittance, zeps : float
        Weight average of geometrical z emittance, [rad-MeV/u].
    xnemittance, xepsn : float
        Weight average of normalized x emittance, [mm-mrad].
    ynemittance, yepsn : float
        Weight average of normalized y emittance, [mm-mrad].
    znemittance, zepsn : float
        Weight average of normalized z emittance, [rad-MeV/u].
    xemittance_all, xeps_all : Array
        Geometrical x emittance of all charge states, [mm-mrad].
    yemittance_all, yeps_all : Array
        Geometrical y emittance of all charge states, [mm-mrad].
    zemittance_all, zeps_all : Array
        Geometrical z emittance of all charge states, [rad-MeV/u].
    xnemittance_all, xepsn_all : Array
        Normalized x emittance of all charge states, [mm-mrad].
    ynemittance_all, yepsn_all : Array
        Normalized y emittance of all charge states, [mm-mrad].
    znemittance_all, zepsn_all : Array
        Normalized z emittance of all charge states, [rad-MeV/u].
    xtwiss_beta, xtwsb : float
        Weight average of twiss beta x, [m/rad].
    ytwiss_beta, ytwsb : float
        Weight average of twiss beta y, [m/rad].
    ztwiss_beta, ztwsb : float
        Weight average of twiss beta z, [rad/MeV/u].
    xtwiss_alpha, xtwsa : float
        Weight average of twiss alpha x, [1].
    ytwiss_alpha, ytwsa : float
        Weight average of twiss alpha y, [1].
    ztwiss_alpha, ztwsa : float
        Weight average of twiss alpha z, [1].
    xtwiss_beta_all, xtwsb_all : Array
        Twiss beta x of all charge states, [m/rad].
    ytwiss_beta_all, ytwsb_all : Array
        Twiss beta y of all charge states, [m/rad].
    ztwiss_beta_all, ztwsb_all : Array
        Twiss beta z of all charge states, [rad/MeV/u].
    xtwiss_alpha_all, xtwsa_all : Array
        Twiss alpha x of all charge states, [1].
    ytwiss_alpha_all, ytwsa_all : Array
        Twiss alpha y of all charge states, [1].
    ztwiss_alpha_all, ztwsa_all : Array
        Twiss alpha z of all charge states, [1].
    couple_xy, cxy : float
        Weigth average of normalized x-y coupling term, [1].
    couple_xpy, cxpy : float
        Weigth average of normalized xp-y coupling term, [1].
    couple_xyp, cxyp : float
        Weigth average of normalized x-yp coupling term, [1].
    couple_xpyp, cxpyp : float
        Weigth average of normalized xp-yp coupling term, [1].
    couple_xy_all, cxy_all : Array
        Normalized x-y coupling term of all charge states, [1].
    couple_xpy_all, cxpy_all : Array
        Normalized xp-y coupling term of all charge states, [1].
    couple_xyp_all, cxyp_all : Array
        Normalized x-yp coupling term of all charge states, [1].
    couple_xpyp_all, cxpyp_all : Array
        Normalized xp-yp coupling term of all charge states, [1].

    Examples
    --------
    >>> fm = ModelFlame(lat_file = 'userfile.lat')
    >>> hplot('xrms', 'yrms', machine=fm)
    """
    ldct = {'rms': 'RMS',
            'cen': 'centroid',
            'eps': 'geom. emittance',
            'epsn': 'norm. emittance',
            'twsb': 'twiss beta',
            'twsa': 'twiss alpha'}
    cdct = {'x': 'b', 'y': 'r', 'z': 'g'}
    if 'colors' in kws:
        for key in kws['colors']:
            cdct[key] = kws['colors'][key]
    lsls = ['-', '--', ':', '-.']
    lttc  = bool(kws['lattice']) if 'lattice' in kws else True
    lgnd = bool(kws['legend']) if 'legend' in kws else True
    cs = int(kws['cs']) if 'cs' in kws else 0
    m = kws['machine'] if 'machine' in kws else None
    r = kws['result'] if 'result' in kws else None

    if isinstance(m, str):
        m = ModelFlame(lat_file=m)

    if m is None and r is None:
        _LOGGER.error('Nothing to plot.')
        return None
    elif r is None:
        if isinstance(m, Machine):
            s = m.allocState({})
            r = m.propagate(s, observe = range(len(m)))
        elif isinstance(m, ModelFlame):
            r, s = m.run(monitor = 'all')
            if r[0][0] != 0:
                r = [(0, m.bmstate)] + r
            m = m.machine
        else:
            _LOGGER.error('Unsuported type of machine.')
            return None
    else:
        if isinstance(m, Machine):
            pass
        elif isinstance(m, ModelFlame):
            m = m.machine
        else:
            _LOGGER.error('Unsuported type of machine.')
            return None

    keys = ['pos'] + list(args)
    d = collect_data(r, *keys)
    for key in keys:
        if len(d[key].shape) == 2:
            d[key] = d[key][:, cs]

    if lttc:
        if len(args) != 0:
            ymax = -1e256
            ymin = 1e256
            for key in args:
                ymax = max([ymax, max(d[key])])
                ymin = min([ymin, min(d[key])])
            ydif = ymax - ymin
            yscl = ydif
            if ydif == 0.0:
                ydif = ymax*0.1 if ymax != 0.0 else 0.1
                yscl = ydif*0.2
        else:
            ymin = 0
            ydif = 0
            yscl = 10

        l = PlotLat(m, auto_scaling=False)
        l.generate(ycen=ymin-0.2*ydif, yscl=0.1*yscl, legend=False, option=False)

    cntx = -1
    cnty = -1
    cntz = -1
    cnto = -1

    for key in args:
        if key[0] == 'x':
            color = cdct[key] if key in cdct else cdct['x']
            cnt = cntx = (cntx + 1)%4
            label = 'x' if cnt == 0 else key
        elif key[0] == 'y':
            color = cdct[key] if key in cdct else cdct['y']
            cnt = cnty = (cnty + 1)%4
            label = 'y' if cnt == 0 else key
        elif key[0] == 'z':
            color = cdct[key] if key in cdct else cdct['z']
            cnt = cntz = (cntz + 1)%4
            label = 'z' if cnt == 0 else key
        else:
            color = cdct[key] if key in cdct else None
            cnt = cnto = (cnto + 1)%4
            label = key

        plt.plot(d['pos'], d[key], c=color, ls=lsls[cnt], label = label)

    if lgnd and len(args) != 0:
        plt.legend(loc='best')

    plt.xlabel('z [m]')
    plt.xlim([min(d['pos']), max(d['pos'])])

    if len(args) != 0:
        ykey = args[0]
        if ykey[0] in ['x', 'y', 'z']:
            ylabel = ykey[1:]
        elif ykey[0:4] == 'ref_':
            ylabel = ykey[4:]
        else:
            ylabel = ykey

        ylabel = ldct[ylabel] if ylabel in ldct else ylabel

        try:
            txt = getattr(BeamState, ykey).__doc__
            u0 = txt.index('[')
            u1 = txt.index(']')
            unit = ' ' + txt[u0:u1+1]
        except:
            unit = ''

        plt.ylabel(ylabel + unit)
    else:
        plt.ylim([-2, 2])

class PlotLat:
    """Lattice picture class from FLAME lattice file or FLAME Machine object.

    Parameters
    ----------
    source : str or callable
        File path of the lattic file (str) or FLAME Machine object (callable)
    output : str (None), optional
        Output file name. If defined, the lattice plot is generated automatically.
    auto_scaling : bool (True), optional
        Flag for y-axis scaling by strength of the optical elements
    starting_offset : float (0.0), optional
        Position offset of starting point in the lattice file

    Attributes
    ----------
    types : dict
        Element type list of the lattice. Each element type contains
        on-off 'flag', plotting 'color', and y-axis 'scale'.

    """
    def __init__(self, source, output=None, auto_scaling=False , starting_offset=0.0, **kws):
        self._source = source
        self._auto_scaling = auto_scaling
        self._starting_offset = starting_offset

        if type(self._source) == str:
            with open(self._source, 'rb') as lat:
                self.M = Machine(lat)

        elif type(self._source) == Machine:
            self.M = self._source

        else:
            raise ValueError('source must be a file path of .lat or flame.Machine object')

        self.types = {'rfcavity':   {'flag':True, 'name':'rfcavity', 'color':'orange', 'scale':0.0},
                      'solenoid':   {'flag':True, 'name':'solenoid', 'color':'red',    'scale':0.0},
                      'quadrupole': {'flag':True, 'name':'quad',     'color':'purple', 'scale':0.0},
                      'sextupole':  {'flag':True, 'name':'sext',     'color':'navy',   'scale':0.0},
                      'sbend':      {'flag':True, 'name':'bend',     'color':'green',  'scale':0.0},
                      'equad':      {'flag':True, 'name':'e-quad',   'color':'blue',   'scale':0.0},
                      'edipole':    {'flag':True, 'name':'e-dipole', 'color':'lime',   'scale':0.0},
                      'bpm':        {'flag':True, 'name':'bpm',      'color':'m',      'scale':0.0},
                      'orbtrim':    {'flag':True, 'name':'corr',     'color':'black',  'scale':0.0},
                      'stripper':   {'flag':True, 'name':'stripper', 'color':'y',      'scale':0.0},
                      'marker':     {'flag':True, 'name':'pm',       'color':'c',      'scale':0.0}
                      }

        if self._auto_scaling:
            for i in range(len(self.M)):
                elem = self.M.conf(i)
                if elem['type'] in self.types.keys():
                    prv_scl = self.types[elem['type']]['scale']
                    tmp_scl = np.abs(self._get_scl(elem))

                    self.types[elem['type']]['scale'] = max(prv_scl,tmp_scl)

        if isinstance(output, str):
            self.generate(legend=True, option=True)
            self.output(window=True, fname=output, **kws)

    def _get_scl(self, elem):
        """Get arbital strength of the optical element.
        """

        scl = 0.0

        if elem['type'] == 'rfcavity':
            scl = elem['scl_fac']*np.cos(2.0*np.pi*elem['phi']/360.0)
        elif elem['type'] == 'solenoid':
            scl = elem['B']
        elif elem['type'] == 'quadrupole':
            scl = elem['B2'] if 'B2' in elem else 1.0
        elif elem['type'] == 'sextupole':
            scl = elem['B3']
        elif elem['type'] == 'sbend':
            scl = elem['phi']
        elif elem['type'] == 'equad':
            scl = elem['V']/elem['radius']**2.0
        elif elem['type'] == 'edipole':
            scl = elem['phi']

        return scl

    def generate(self, start=None, end=None, xlim=None, ycen=0.0, yscl=1.0, aspect=5.0,
                 legend=True, option=True, axes = None):
        """Generate matplotlib Axes class object from lattice file.

        Parameters
        ----------
        start : int
            Index of the lattice start.
        end : int
            Index of the lattice end.
        xlim : list[2], optinal
            Plot range of the lattice.
        ycen : float (0.0)
            Vertical offset of the plot.
        yscl : float (1.0)
            Plot scale of the element size
        aspect : float (5.0), optional
            Aspect ratio of the picture.
        legend : bool
            Add legend for the elements.
        option : bool
            Add optional setting for the plot.

        Attributes
        ----------
        axes : callable
            Axes class object of matplotlib.

        total_length : float
            Total length of the lattice.

        """

        if axes is None:
            self._fig = plt.figure()
            self.axes = self._fig.add_subplot(111)
        else:
            self.axes = axes

        pos = self._starting_offset
        bp = ycen
        indexes = range(len(self.M))[start:end]
        foundelm = []

        for i in indexes:
            elem = self.M.conf(i)

            try:
                dL = elem['L']
            except:
                dL = 0.0

            if elem['type'] in self.types.keys():
                info = self.types[elem['type']]

                if foundelm.count(elem['type']) == 0:
                    foundelm.append(elem['type'])
                    if legend and info['flag']:
                        self.axes.fill_between([0,0],[0,0],[0,0], color=info['color'], label=info['name'])

                if info['flag']:
                    if dL != 0.0:
                        bpp = bp
                        if info['scale'] != 0.0:
                            ht = yscl*self._get_scl(elem)/info['scale'] + 0.05
                        else:
                            ht = yscl*np.sign(self._get_scl(elem))

                        if elem['type'] == 'rfcavity' or elem['type'] == 'solenoid':
                            bpp = bp-yscl*0.7
                            ht = yscl*2.0*0.7

                        self.axes.add_patch(ptc.Rectangle((pos, bpp), dL, ht,
                                                           edgecolor='none',facecolor=info['color']))
                    else:
                        self.axes.add_line(lin.Line2D([pos,pos],[-yscl*0.3+bp, yscl*0.3+bp],color=info['color']))

            pos += dL

        self.total_length = pos
        self.axes.add_line(lin.Line2D([0.0, pos], [bp,bp], color='gray', zorder=-5))

        if len(foundelm) <= 4 :
            ncol = 4
        elif len(foundelm) <= 6:
            ncol = 3
        elif len(foundelm) <= 8 :
            ncol = 4
        elif len(foundelm) <= 9:
            ncol = 3
        else :
            ncol = 4

        if option:
            if xlim != None:
                self.axes.set_xlim(xlim)
                ancscl = xlim[1]-xlim[0]
            else :
                self.axes.set_xlim((0.0,pos))
                ancscl = pos

            if legend:
                self.axes.legend(ncol=ncol, loc=10, bbox_to_anchor=(0.5, -0.2*ancscl))

            self.axes.set_aspect(aspect)
            self.axes.set_ylim((-1.0,1.0))
            self.axes.set_yticks(())
            self.axes.grid()
            self.axes.spines['right'].set_visible(False)
            self.axes.spines['left'].set_visible(False)
            self.axes.spines['top'].set_visible(False)
            self.axes.spines['bottom'].set_visible(False)
            plt.tight_layout()

    def output(self,window=True,fname=None,**kws):
        """Output the lattice picture to window and/or file.

        Parameters
        ----------
        window : bool (True)
            Output flag to the window.
        fname : str, optional
            File path of the output picutre.
        **kwargs :
            The same kwargs as pyplot.savefig() are available.
        """

        if type(fname) == str :
            plt.savefig(fname, **kws)

        if window: plt.show()

if __name__ == "__main__":
    import sys

    try:
        lat = sys.argv[1]
    except:
        raise ValueError('First argument must be a lattice file.')

    try:
        out = sys.argv[2]
    except:
        raise ValueError('Second argument must be a output file name.')

    pl = PlotLat(lat, output=out)


#!/usr/bin/env python3

from .metric import ReferencedMetric
import numpy as np
from rouge_score import rouge_scorer, scoring

class ROUGE(ReferencedMetric):
    """ROUGE uses Google implementation (https://github.com/google-research/google-research/tree/master/rouge)
    but adds own implementation of multi-ref jackknifing.
    The Google implementation should be identical to Rouge-155 (except tokenization?),
    the jackknifing follows the description of the ROUGE paper.
    """

    def compute(self, cache, predictions, references):
        rouge_types = ["rouge1", "rouge2", "rougeL", "rougeLsum"]
        rouge = rouge_scorer.RougeScorer(rouge_types=rouge_types, use_stemmer=True)
        scores = {}
        # TODO expecting pretokenized data, do we want to imitate Rouge-155 tokenizer somehow?
        for refs, pred, pred_id in zip(
            references.whitespace_tokenized,
            predictions.whitespace_tokenized,
            predictions.ids,
        ):
            # ROUGE multi-ref jackknifing
            if len(refs) > 1:
                cur_scores = [rouge.score(ref, pred) for ref in refs]

                # get best score for all leave-one-out sets
                best_scores = []
                for leave in range(len(refs)):
                    cur_scores_leave_one = [
                        cur_scores[s] for s in range(len(refs)) if s != leave
                    ]
                    best_scores.append(
                        {
                            rouge_type: max(
                                [s[rouge_type] for s in cur_scores_leave_one],
                                key=lambda s: s.fmeasure,
                            )
                            for rouge_type in rouge_types
                        }
                    )

                # average the leave-one-out bests to produce the final score
                score = {
                    rouge_type: scoring.Score(
                        np.mean([b[rouge_type].precision for b in best_scores]),
                        np.mean([b[rouge_type].recall for b in best_scores]),
                        np.mean([b[rouge_type].fmeasure for b in best_scores]),
                    )
                    for rouge_type in rouge_types
                }
            else:
                score = rouge.score(refs[0], pred)

            # convert the named tuples to plain nested dicts
            score = {
                rouge_type: {
                    "precision": score[rouge_type].precision,
                    "recall": score[rouge_type].recall,
                    "fmeasure": score[rouge_type].fmeasure,
                }
                for rouge_type in rouge_types
            }
            # Write to cache if not None.
            if cache is not None:
                cache_key = (self.__class__.__name__, predictions.filename, pred_id)
                cache[cache_key] = score
            scores[pred_id] = score

        return scores
        # return result


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@date: 2021-Today
@author: marie-christin.wimmler@tu-dresden.de
"""

from TreeModelLib.BelowgroundCompetition.OGSLargeScale3D import OGSLargeScale3D
from ProjectLib.Logger import method_logger
import numpy as np

# Child class of OGSLargeScale3DExternal to use external time stepping
# E.g. to run MANGA as OGS boundary condition
# The concept needs an array with cumulated cell salinity and
# the number of calls for each cell. It returns an array describing water
# withdrawal in each cell as rate in kg per sec per cell volume.
# MRO: OGSLargeScale3DExternal, OGSLargeScale3D, TreeModel, object
class OGSLargeScale3DExternal(OGSLargeScale3D):
    @method_logger
    def __init__(self, args):
        super().__init__(args)

    # allow external communication
    def getOGSAccessible(self):
        return True

    ## This functions prepares the next timestep for the competition
    #  concept. In the OGS concept, information on t_ini and t_end is stored.
    #  Additionally, arrays are prepared to store information on water uptake
    #  of the participating trees. Moreover, the ogs-prj-file for the next
    #  time step is updated and saved in the ogs-project folder.
    #  @param t_ini: initial time of next time step
    #  @param t_end: end time of next time step
    @method_logger
    def prepareNextTimeStep(self, t_ini, t_end):
        self.n_trees = 0

        # Arrays with length 'no. of trees'
        self._total_resistance = []

        self._psi_leaf = np.empty(0)
        self._psi_height = np.empty(0)
        self._psi_osmo = np.empty(0)
        self._rcrown = []
        self._hstem = []

        super().prepareOGSparameters()

    ## Before being able to calculate the resources, all tree enteties need
    #  to be added with their current implementation for the next time step.
    #  Here, in the OGS case, each tree is represented by a contribution to
    #  python source terms in OGS.
    #  @param tree
    @method_logger
    def addTree(self, tree):
        x, y = tree.getPosition()
        geometry = tree.getGeometry()
        parameter = tree.getParameter()

        # Cells affected by tree water uptake
        super().addCellCharateristics(x, y)

        # Resistances against water flow in tree
        # Calculate total tree resistance
        total_resistance = super().totalTreeResistance(parameter, geometry)
        self._total_resistance.append(total_resistance)

        # Water potentials
        self._psi_leaf = np.concatenate(
            (self._psi_leaf, [(parameter["leaf_water_potential"])]))
        self._psi_height = np.concatenate(
            (self._psi_height,
             [-(2 * geometry["r_crown"] + geometry["h_stem"]) * 9810]))

    ## This function updates and returns BelowgroundResources in the current
    #  time step. For each tree a reduction factor is calculated which is
    #  defined as: resource uptake at zero salinity/ real resource uptake.
    @method_logger
    def calculateBelowgroundResources(self):
        # Number of trees
        self.no_trees = len(self._total_resistance)
        # Salinity below each tree
        self._tree_salinity = np.empty(self.no_trees)
        if self.no_trees <= 0:
            print("WARNING: All trees are dead.")

        # Calculate salinity (and psi_osmo) below tree
        super().calculateTreeSalinity()

        self._tree_water_uptake = -(self._psi_leaf - self._psi_height -
                                   self._psi_osmo) / \
                                 self._total_resistance / np.pi * 1000  # kg/s
        self.belowground_resources = 1 - (self._psi_osmo /
                                          (self._psi_leaf - self._psi_height))

        # Calculate contribution per cell
        super().calculateCompleteTreeContribution()

    ## Setter for external information
    # This function sets the parameters 'cumsum_salinity' and 'calls_per_cell',
    # which contain information about the cumulated salinity in each cell and
    # the number of calls, calculated by OGS
    @method_logger
    def setExternalInformation(self, **args):
        # information about cell salinity from OGS
        self.cumsum_salinity = args["cumsum_salinity"]
        self.calls_per_cell = args["calls_per_cell"]
        self._salinity = self.cumsum_salinity / self.calls_per_cell

    ## Getter for external information
    # This function returns the estimated water withdrawal in each cell
    # as rate (kg per sec per cell volume)
    def getExternalInformation(self):
        return self._tree_contribution_per_cell


# -*- coding: utf-8 -*-
import pickle
from keras.models import load_model
from .data_processing import data_process
import numpy as np
import os
from keras.preprocessing.sequence import pad_sequences

localpath = os.path.dirname(__file__)

class Predictor(object):
    def __init__(self, num_words=50000, max_len=500,
                 path_tokenizer=os.path.join(localpath, 'model/tokenizer_50000_cleaned.pkl'),
                 path_accusation=None,
                 path_relative_articles=os.path.join(localpath,
                                                     'model/textcnn_relevant_articles_cleaned_token_50000_pad_500_filter_512_hidden_1000_epoch_3_accu_80_f1_80.h5')):
        self.num_words = num_words
        self.max_len = max_len
        self.path_accusation = path_accusation
        self.path_relative_articles = path_relative_articles
        self.batch_size = 500
        self.content_process = data_process()
        self.path_tokenizer = path_tokenizer
        self.model_relative_articles = load_model(path_relative_articles)

    def predict(self, content):
        content_process = self.content_process
        content_seg = content_process.segmentation(content, cut=True, word_len=2, replace_money_value=True,
                                                   stopword=True)

        with open(self.path_tokenizer, mode='rb') as f:
            tokenizer = pickle.load(f)

        content_process.text2num(content_seg, tokenizer=tokenizer)
        content_seg_num_sequence = content_process.num_sequence
        content_fact_pad_seq = pad_sequences(content_seg_num_sequence, maxlen=self.max_len, padding='post')
        content_fact_pad_seq = np.array(content_fact_pad_seq)

        model_relative_articles = self.model_relative_articles
        relative_articles = model_relative_articles.predict(content_fact_pad_seq)

        def transform(x):
            n = len(x)
            x_return = np.arange(1, n + 1)[x > 0.5].tolist()
            if len(x_return) == 0:
                x_return = np.arange(1, n + 1)[x == x.max()].tolist()
            return x_return

        result = []
        for i in range(0, len(content)):
            result.append({
                "accusation": [None],
                "imprisonment": 0,
                "articles": transform(relative_articles[i])
            })
        return result

# if __name__ == '__main__':
#     content = ['公诉机关起诉指控，被告人张某某秘密窃取他人财物', '锡林浩特市人民检察院指控，被告人杨某某以非法占有为目的，秘密窃取他人财物']
#     predictor = Predictor()
#     m = predictor.predict(content)
#     print(m)


import numpy as np

from privacy_evaluator.attacks.attack import Attack
from privacy_evaluator.classifiers.classifier import Classifier

class SampleAttack(Attack):
    """SampleAttack class."""

    def __init__(
        self,
        target_model: Classifier,
        x_train: np.ndarray,
        y_train: np.ndarray,
        x_test: np.ndarray,
        y_test: np.ndarray,
    ):
        """Initializes a SampleAttack class.

        :param target_model: Target model to be attacked.
        :param x_train: Data that was used to train the target model.
        :param y_train: Labels for the data that was used to train the target model.
        :param x_test: Data that was not used to train the target model.
        :param y_test: Labels for the data that was not used to train the target model.
        """
        super().__init__(target_model, x_train, y_train, x_test, y_test)

    def attack(self, x: np.ndarray, y: np.ndarray, *args, **kwargs) -> np.ndarray:
        """Performs the attack on the target model.

        :param x: Input data to attack.
        :param y: True labels for x.
        :param args: Arguments of the attack.
        :param kwargs: Keyword arguments of the attack.
        :return: An array holding the result of the attack.
        """
        return np.ones(10)