from django.shortcuts import render, redirect
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.http import HttpResponseRedirect, HttpResponse, Http404
from django.urls import reverse

from django.contrib import messages
import datetime

from django.db.models import Q
from .models import *

@login_required
def index(request):
    service_posts = ServicePost.objects.filter().order_by('-created_datetime')[:10]

    return render(request, 'service/index.html', {
        'navigation': 'service',
        'service_posts': service_posts,
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def explore(request):
    service_posts = ServicePost.objects.filter().order_by('-created_datetime')
    return render(request, 'service/explore.html', {
        'navigation': 'service',
        'service_posts': service_posts,
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def search(request):

    if request.method == 'GET':
        search_text = request.GET.get('service_search_text')

        title_matches = ServicePost.objects.filter(
            Q(title__icontains=search_text) | Q(service__name__icontains=search_text)
        )

        description_matches = ServicePost.objects.filter(
            Q(description__icontains=search_text) | Q(service__info__icontains=search_text)
        )

        location_matches = ServicePost.objects.filter(
            Q(service__location__icontains=search_text)
        )

        tag_matches = ServicePost.objects.filter(

        ).filter(service__tags__label__icontains=search_text)

    return render(request, 'service/search.html', {
        'search_text': search_text,
        'navigation': 'service',
        'title_matches': title_matches,
        'description_matches': description_matches,
        'location_matches': location_matches,
        'tag_matches': tag_matches,
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def create(request):
    if request.method == 'POST':
        # Creating new service and saving it
        new_service = Service()
        new_service.name = request.POST.get('service-name')
        new_service.info = request.POST.get('service-description')
        new_service.start_datetime = request.POST.get('service-start-datetime')
        new_service.end_datetime = request.POST.get('service-end-datetime')
        new_service.location = request.POST.get('service-location')

        new_service.save()
        tags = request.POST.getlist('service-tags')
        for tag_name in tags:
            tag, created = Tag.objects.get_or_create(
                label = tag_name
            )

            new_service.tags.add(tag)

        new_service.save()

        # Creating new service post record and saving it
        new_post = ServicePost()
        new_post.title = request.POST.get('post-title')
        new_post.description= request.POST.get('post-description')
        new_post.is_request_post = (request.POST['is_request_post'] == '1')
        new_post.owner = User.objects.get(username=request.user.username)
        new_post.created_datetime = datetime.datetime.now()
        new_post.service = new_service

        new_post.save()

        notification = Notification()
        notification.owner = request.user
        notification.created_datetime = datetime.datetime.now()
        notification.text = "You created a new service post titled '" + request.POST.get('post-title') + "'"
        notification.target_url = ""
        notification.save()

        messages.success(request, "Post Published!", extra_tags='message-head')

        messages.success(request, 'Hey ' + request.user.first_name + '! We have recorded and published your post!', extra_tags="message-body")

        return HttpResponseRedirect(reverse('service:create'))

    return render(request, 'service/create.html', {
        'navigation': 'service',
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def detail(request, post_id):
    try:
        post = ServicePost.objects.get(pk=post_id)
        comments = ServicePostComment.objects.filter(post=post).order_by('-created_datetime')

    except ServicePost.DoesNotExist:
        raise Http404('Post does not exist')

    return render(request, 'service/details.html', {
        'navigation': 'service',
        'post': post,
        'comments': comments,
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def update(request, post_id):
    try:
        post = ServicePost.objects.get(pk=post_id)
    except ServicePost.DoesNotExist:
        raise Http404('Post does not exist')

    if request.method == 'POST':

        post.service.name = request.POST.get('service-name')
        post.service.info = request.POST.get('service-description')
        post.service.start_datetime = request.POST.get('service-start-datetime')
        post.service.end_datetime = request.POST.get('service-end-datetime')
        post.service.location = request.POST.get('service-location')

        post.service.save()
        post.save()

        post.service.tags.clear()
        tags = request.POST.getlist('service-tags')
        for tag_name in tags:
            tag, created = Tag.objects.get_or_create(
                label = tag_name
            )

            post.service.tags.add(tag)

        post.service.save()
        post.save()

        post.title = request.POST.get('post-title')
        post.description= request.POST.get('post-description')
        post.is_request_post = (request.POST['is_request_post'] == '1')
        post.modified_datetime = datetime.datetime.now()

        post.save()

        notification = Notification()
        notification.owner = request.user
        notification.created_datetime = datetime.datetime.now()
        notification.text = "You updated your service post information. The post is titled '" + request.POST.get('post-title') + "'"
        notification.target_url = ""
        notification.save()

        messages.success(request, "Post Updated!", extra_tags='message-head')

        messages.success(request, 'Hey ' + request.user.first_name + '! We have updated your post!', extra_tags="message-body")

        return HttpResponseRedirect(reverse('service:update', kwargs={'post_id': post.id}))

    return render(request, 'service/update.html', {
        'navigation': 'service',
        'post': post,
        'unread_notifications_count': Notification.objects.filter(is_read=False,owner=request.user).count()
    })

@login_required
def close(request, post_id):
    try:
        post = ServicePost.objects.get(pk=post_id)
    except ServicePost.DoesNotExist:
        raise Http404('Post does not exist')

    if request.user != post.owner:
        messages.warning(request, "Unauthorized Attempt", extra_tags='message-head')
        messages.warning(request, "You're not authorized to perform this operation. Only the post owner can close the post.", extra_tags='message-body')

    else:
        post.closed_datetime = datetime.datetime.now()
        post.save()

        notification = Notification()
        notification.owner = request.user
        notification.created_datetime = datetime.datetime.now()
        notification.text = "You recently closed a service post titled '" + post.title + "'"
        notification.target_url = ""
        notification.save()

        messages.success(request, "Post closed", extra_tags='message-head info')
        messages.warning(request, "You've recently closed a service post. You will no require receive comments in that post, although the post is still visible unless you delete it!", extra_tags='message-body')

    return HttpResponseRedirect(reverse('service:detail', kwargs={'post_id': post.id}))

@login_required
def delete(request, post_id):
    try:
        post = ServicePost.objects.get(pk=post_id)
    except ServicePost.DoesNotExist:
        raise Http404('Post does not exist')

    if request.user != post.owner:
        messages.warning(request, "Unauthorized Attempt", extra_tags='message-head error')
        messages.warning(request, "You're not authorized to perform this operation. Only the post owner can delete the post.", extra_tags='message-body')

    else:
        post.service.delete()
        post.delete()

        notification = Notification()
        notification.owner = request.user
        notification.created_datetime = datetime.datetime.now()
        notification.text = "You deleted your service post titled '" + post.title + "'"
        notification.target_url = ""
        notification.save()

        messages.success(request, "Post Deleted", extra_tags='message-head info')
        messages.warning(request, "You've successfully deleted a service post. You will no longer be able to access the post!", extra_tags='message-body')

    return HttpResponseRedirect(reverse('service:index'))

@login_required
def create_comment(request, post_id):
    try:
        post = ServicePost.objects.get(pk=post_id)
    except ServicePost.DoesNotExist:
        raise Http404('Post does not exist')

    if request.method == 'POST':
        comment_text = request.POST.get('comment')

        new_comment = ServicePostComment()
        new_comment.text = comment_text
        new_comment.created_datetime = datetime.datetime.now()

        new_comment.owner = request.user
        new_comment.post = post

        new_comment.is_private = False
        if request.POST.get('is-private'):
            new_comment.is_private = True

        new_comment.save()

        notification = Notification()
        notification.owner = post.owner
        notification.created_datetime = datetime.datetime.now()
        notification.text = request.user.first_name + " " + request.user.last_name + " commented in your service post!"
        notification.target_url = ""
        notification.save()

        messages.success(request, "New Comment Added!", extra_tags='message-head info')
        messages.warning(request, "Your comment has been added to the post!", extra_tags='message-body')

    return HttpResponseRedirect(reverse('service:detail', kwargs={ 'post_id': post.id }))


from __future__ import division

from util.inception_resnet_v1 import *
from tensorflow.python.ops import data_flow_ops
from util.file_reader import *
from util.progress import *
from datetime import datetime

class DualReferenceFR(object):
    def __init__(self):
        # data directories
        # self.data_dir = '/home/bingzhang/Documents/Dataset/CACD/CACD2000'
        # self.data_info = '/home/bingzhang/Documents/Dataset/CACD/celenew.mat'
        # self.val_dir = '/home/bingzhang/Documents/Dataset/ZID/LFW/lfw'

        self.data_dir = '/scratch/BingZhang/dataset/CACD2000_Cropped'
        self.data_info = '/scratch/BingZhang/dataset/CACD2000/celenew2.mat'
        self.val_dir = '/scratch/BingZhang/lfw_250/'
        self.val_list = './data/val_list.txt'
        # model directory
        self.log_dir = '/scratch/BingZhang/logs_all_in_one/drfr_dual_train'
        self.model_dir = '/scratch/BingZhang/models_all_in_one/DRFRQDualTrain0.006lrbatch30-Model'
        # image size
        self.image_height = 250
        self.image_width = 250
        self.image_channel = 3
        # net parameters
        self.step = 0
        self.learning_rate = 0.006
        self.batch_size = 30  # must be a multiple of 3
        self.feature_dim = 1024
        self.embedding_size = 128
        self.max_epoch = 10000
        self.delta = 0.2
        self.nof_sampled_id = 45
        self.nof_images_per_id = 20
        self.id_sampled_examples = self.nof_images_per_id * self.nof_sampled_id
        self.nof_sampled_age = 20
        self.nof_images_per_age = 45
        self.age_sampled_examples = self.nof_images_per_age * self.nof_sampled_age
        self.val_size = 144
        # placeholder
        self.path_placeholder = tf.placeholder(tf.string, [None, 3], name='paths')
        self.label_placeholder = tf.placeholder(tf.int64, [None, 3], name='labels')
        self.batch_size_placeholder = tf.placeholder(tf.int32, name='batch_size')

        # for id
        self.id_affinity_watch = tf.placeholder(tf.float32,
                                                [None, self.id_sampled_examples, self.id_sampled_examples, 1])
        self.id_affinity_watch_binarized = tf.placeholder(tf.float32,
                                                          [None, self.id_sampled_examples, self.id_sampled_examples, 1])
        self.affinity_on_val = tf.placeholder(tf.float32, [None, self.val_size, self.val_size, 1])
        self.nof_id_triplets_placeholder = tf.placeholder(tf.int16, name='nof_id_triplets')

        # for age
        self.age_affinity_watch = tf.placeholder(tf.float32,
                                                 [None, self.age_sampled_examples, self.age_sampled_examples, 1])
        self.age_affinity_watch_binarized = tf.placeholder(tf.float32,
                                                           [None, self.age_sampled_examples, self.age_sampled_examples,
                                                            1])
        self.nof_age_triplets_placeholder = tf.placeholder(tf.int16, name='nof_age_triplets')
        # ops
        self.input_queue = data_flow_ops.FIFOQueue(capacity=1000000, dtypes=[tf.string, tf.int64], shapes=[(3,), (3,)])
        self.enqueue_op = self.input_queue.enqueue_many([self.path_placeholder, self.label_placeholder])

        nof_process_threads = 4
        images_and_labels = []
        for _ in range(nof_process_threads):
            file_paths, labels = self.input_queue.dequeue()
            images = []
            for file_path in tf.unstack(file_paths):
                file_content = tf.read_file(file_path)
                image = tf.image.decode_png(file_content)
                image.set_shape((self.image_height, self.image_width, self.image_channel))
                images.append(tf.image.per_image_standardization(image))
            images_and_labels.append([images, labels])
        self.image_batch, self.label_batch = tf.train.batch_join(images_and_labels,
                                                                 batch_size=self.batch_size_placeholder,
                                                                 enqueue_many=True,
                                                                 capacity=nof_process_threads * self.batch_size,
                                                                 shapes=[
                                                                     (self.image_height, self.image_width,
                                                                      self.image_channel), ()],
                                                                 allow_smaller_final_batch=True)

        self.embedding = self.net_forward(self.image_batch)
        self.id_embedding = self.get_id_embeddings(self.embedding)
        self.id_loss = self.get_triplet_loss_v2(self.id_embedding)
        self.age_embedding = self.get_age_embeddings(self.embedding)
        self.age_loss = self.get_triplet_loss_v2(self.age_embedding)
        self.age_summary_op, self.id_summary_op ,self.average_op = self.build_summary()
        with tf.control_dependencies([self.average_op]):
            self.id_opt = tf.train.MomentumOptimizer(learning_rate=self.learning_rate, momentum=0.9,
                                                     use_nesterov=True).minimize(self.id_loss)
            self.age_opt = tf.train.MomentumOptimizer(learning_rate=self.learning_rate, momentum=0.9,
                                                      use_nesterov=True).minimize(self.age_loss)

    def net_forward(self, image_batch):
        net, _ = inference(image_batch, keep_probability=1.0, bottleneck_layer_size=128, phase_train=True,
                           weight_decay=0.0, reuse=None)
        feature = slim.fully_connected(net, self.feature_dim, activation_fn=None,
                                       weights_initializer=tf.truncated_normal_initializer(stddev=0.1),
                                       weights_regularizer=slim.l2_regularizer(0.0))
        return feature

    def get_id_embeddings(self, feature):
        with tf.variable_scope('id_embedding'):
            id_embeddings = slim.fully_connected(feature, self.embedding_size, activation_fn=None,
                                                 weights_initializer=tf.truncated_normal_initializer(stddev=0.1),
                                                 weights_regularizer=slim.l2_regularizer(0.0), scope='id_embedding')
            weights = slim.get_model_variables('id_embedding')[0]
            bias = slim.get_model_variables('id_embedding')[1]
            variable_summaries(weights, 'weight')
            variable_summaries(bias, 'bias')
            id_embeddings = tf.nn.l2_normalize(id_embeddings, dim=1, epsilon=1e-12, name='id_embeddings')
        return id_embeddings

    def get_age_embeddings(self, feature):
        with tf.variable_scope('age_embedding'):
            age_embeddings = slim.fully_connected(feature, self.embedding_size, activation_fn=None,
                                                  weights_initializer=tf.truncated_normal_initializer(stddev=0.1),
                                                  weights_regularizer=slim.l2_regularizer(0.0), scope='age_embedding')
            weights = slim.get_model_variables('age_embedding')[0]
            bias = slim.get_model_variables('age_embedding')[1]
            variable_summaries(weights, 'weight')
            variable_summaries(bias, 'bias')
            age_embeddings = tf.nn.l2_normalize(age_embeddings, dim=1, epsilon=1e-12, name='id_embeddings')
        return age_embeddings

    def get_triplet_loss(self, embeddings):
        anchor = embeddings[0:self.batch_size:3][:]
        positive = embeddings[1:self.batch_size:3][:]
        negative = embeddings[2:self.batch_size:3][:]

        pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), 1)
        neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), 1)

        with tf.name_scope('distances'):
            tf.summary.histogram('positive', pos_dist)
            tf.summary.histogram('negative', neg_dist)

        basic_loss = tf.add(tf.subtract(pos_dist, neg_dist), self.delta)
        loss = tf.reduce_mean(tf.maximum(basic_loss, 0.0), 0)
        return loss

    def get_triplet_loss_v2(self, embeddings):
        with tf.name_scope('loss_012'):
            anchor = embeddings[0:self.batch_size:3][:]
            positive = embeddings[1:self.batch_size:3][:]
            negative = embeddings[2:self.batch_size:3][:]
            pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), 1)
            neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), 1)

            with tf.name_scope('distances'):
                tf.summary.histogram('positive_012', pos_dist)
                tf.summary.histogram('negative_012', neg_dist)

            basic_loss = tf.add(tf.subtract(pos_dist, neg_dist), self.delta)
            loss_012 = tf.reduce_mean(tf.maximum(basic_loss, 0.0), 0)

        with tf.name_scope('loss_102'):
            anchor = embeddings[1:self.batch_size:3][:]
            positive = embeddings[0:self.batch_size:3][:]
            negative = embeddings[2:self.batch_size:3][:]
            pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), 1)
            neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), 1)

            with tf.name_scope('distances'):
                tf.summary.histogram('positive_102', pos_dist)
                tf.summary.histogram('negative_102', neg_dist)

            basic_loss = tf.add(tf.subtract(pos_dist, neg_dist), self.delta)
            loss_102 = tf.reduce_mean(tf.maximum(basic_loss, 0.0), 0)
        return loss_012 + loss_102

    def build_summary(self):
        with tf.name_scope('affinity'):
            tf.summary.image('id_original', self.id_affinity_watch)
            tf.summary.image('id_binarized', self.id_affinity_watch_binarized)
            tf.summary.image('age_original', self.age_affinity_watch)
            tf.summary.image('age_binarized', self.age_affinity_watch_binarized)
            tf.summary.image('val', self.affinity_on_val)
        with tf.name_scope('loss'):
            tf.summary.scalar('id_loss', self.id_loss)
            tf.summary.scalar('age_loss', self.age_loss)

            average = tf.train.ExponentialMovingAverage(0.9)
            average_op = average.apply([self.id_loss, self.age_loss])
            tf.summary.scalar('id_loss_aver', average.average(self.id_loss))
            tf.summary.scalar('age_loss_aver', average.average(self.age_loss))
        with tf.name_scope('embeddings'):
            tf.summary.histogram('id_embeddings', self.id_embedding)
            tf.summary.histogram('age_embeddings', self.age_embedding)
        with tf.name_scope('nof_triplet'):
            tf.summary.scalar('nof_id_triplet', self.nof_id_triplets_placeholder)
            tf.summary.scalar('nof_age_triplet', self.nof_age_triplets_placeholder)

        var = tf.get_collection(tf.GraphKeys.SUMMARIES)
        age_sum = [v for v in var if str(v).__contains__('age')]
        id_sum = [v for v in var if str(v).__contains__('id')]
        others = [v for v in var if not str(v).__contains__('age') and not str(v).__contains__('id')]
        return tf.summary.merge([age_sum,others]),tf.summary.merge([id_sum,others]), average_op

    # def test_module(self):
    #     CACD = FileReader(self.data_dir,self.data_info,reproducible=True,contain_val=False)
    #     path_array, label_array = CACD.select_identity_path(self.nof_sampled_id, self.nof_images_per_id)
    #     path_array = np.reshape(path_array,(-1,1))
    #
    #     label_array = np.reshape(np.arange(self.nof_sampled_id*self.nof_images_per_id),(-1,1))
    #     # label_array = np.reshape(label_array,(-1,1))
    #     for step in range(4):
    #         print('before %d enqueue, size is %d' % (step, self.sess.run(self.input_queue.size())))
    #         self.sess.run(self.enqueue_op,
    #                       feed_dict={self.path_place_holder: path_array, self.label_place_holder: label_array})
    #         sum,label= self.sess.run([self.summary_op,self.label_batch])
    #         # sio.savemat('test.mat',{'im':image})
    #         # sum = self.sess.run(self.summary_op)
    #         print(label)
    #         sum,label= self.sess.run([self.summary_op,self.label_batch])
    #         print(label)
    #         print('after dequeue, size is %d' %self.sess.run(self.input_queue.size()))
    #         self.summary_writer.add_summary(sum, step)

    def train(self):
        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        coord = tf.train.Coordinator()
        tf.train.start_queue_runners(coord=coord, sess=sess)

        CACD = FileReader(self.data_dir, self.data_info, val_data_dir=self.val_dir, val_list=self.val_list,
                          reproducible=True, contain_val=True)
        summary_writer = tf.summary.FileWriter(
            os.path.join(self.log_dir, datetime.now().isoformat()), sess.graph)
        var = tf.trainable_variables()
        var = [v for v in var if str(v).__contains__('Inception')]
        saver = tf.train.Saver(var)
        saver.restore(sess, '/scratch/BingZhang/facenet4drfr/model/20170512-110547/model-20170512-110547.ckpt-250000')
        # saver.restore(sess, '/home/bingzhang/Workspace/PycharmProjects/20170512-110547/model-20170512-110547.ckpt-250000')
        aff_val = np.zeros((self.val_size, self.val_size))
        for triplet_selection in range(self.max_epoch):

            # ID Step
            # select some examples to forward propagation
            paths, labels = CACD.select_identity_path(self.nof_sampled_id, self.nof_images_per_id)
            nof_examples = len(paths)
            path_array = np.reshape(paths, (-1, 3))
            label_array = np.reshape(np.arange(nof_examples), (-1, 3))
            embeddings_array = np.zeros(shape=(nof_examples, self.embedding_size))

            # FIFO enqueue
            sess.run(self.enqueue_op,
                     feed_dict={self.path_placeholder: path_array, self.label_placeholder: label_array})

            # forward propagation to get current embeddings
            print('ID STEP\nForward propagation to get current embeddings')
            nof_batches = int(np.ceil(nof_examples / self.batch_size))
            for i in range(nof_batches):
                batch_size = min(nof_examples - i * self.batch_size, self.batch_size)
                emb, label = sess.run([self.id_embedding, self.label_batch],
                                      feed_dict={self.batch_size_placeholder: batch_size})
                embeddings_array[label, :] = emb

            # compute affinity matrix on batch
            aff = []
            for index in range(nof_examples):
                aff.append(np.sum(np.square(embeddings_array[index][:] - embeddings_array), 1))
            aff_binarized = binarize_affinity(aff, self.nof_images_per_id)
            # affinity matrix on

            # select triplets to train
            triplet = select_triplets(embeddings_array, self.nof_sampled_id, self.nof_images_per_id, self.delta)
            triplet_path_array = paths[triplet][:]
            triplet_label_array = labels[triplet][:]
            nof_triplets = len(triplet_path_array)
            print('%d triplets selected' % nof_triplets)

            # FIFO enqueue
            sess.run(self.enqueue_op, feed_dict={self.path_placeholder: triplet_path_array,
                                                 self.label_placeholder: triplet_label_array})
            # train on selected triplets
            nof_batches = int(np.ceil(nof_triplets * 3 / self.batch_size))
            for i in range(nof_batches):
                batch_size = min(nof_triplets * 3 - i * self.batch_size, self.batch_size)
                if self.step % 200 == 0:
                    _sum, loss, _ = sess.run(
                        [self.id_summary_op, self.id_loss, self.id_opt],
                        feed_dict={self.batch_size_placeholder: batch_size,
                                   self.id_affinity_watch: np.reshape(aff, [1, self.id_sampled_examples,
                                                                            self.id_sampled_examples,
                                                                            1]),
                                   self.id_affinity_watch_binarized: np.reshape(aff_binarized, [1,
                                                                                self.id_sampled_examples,
                                                                                self.id_sampled_examples,
                                                                                1]),
                                   self.nof_id_triplets_placeholder: nof_triplets,
                                   self.affinity_on_val: np.reshape(aff_val, [1, self.val_size, self.val_size, 1])})

                    # write in summary
                    summary_writer.add_summary(_sum, self.step)
                else:
                    loss, _ = sess.run([self.id_loss, self.id_opt], feed_dict={self.batch_size_placeholder: batch_size})
                progress(i + 1, nof_batches, str(triplet_selection) + 'th ID Epoch',
                         'Batches loss:' + str(loss))  # a command progress bar to watch training progress
                self.step += 1
                # save model
                if self.step % 200000 == 0:
                    saver.save(sess, self.model_dir, global_step=self.step)

            for inner_iter in range(4):
                # AGE Step
                # select some examples to forward propagation
                paths, labels = CACD.select_age_path(self.nof_sampled_age, self.nof_images_per_age)
                nof_examples = len(paths)
                path_array = np.reshape(paths, (-1, 3))
                label_array = np.reshape(np.arange(nof_examples), (-1, 3))
                embeddings_array = np.zeros(shape=(nof_examples, self.embedding_size))

                # FIFO enqueue
                sess.run(self.enqueue_op,
                         feed_dict={self.path_placeholder: path_array, self.label_placeholder: label_array})

                # forward propagation to get current embeddings
                print('AGE STEP\nForward propagation to get current embeddings')
                nof_batches = int(np.ceil(nof_examples / self.batch_size))
                for i in range(nof_batches):
                    batch_size = min(nof_examples - i * self.batch_size, self.batch_size)
                    emb, label = sess.run([self.age_embedding, self.label_batch],
                                          feed_dict={self.batch_size_placeholder: batch_size})
                    embeddings_array[label, :] = emb

                # compute affinity matrix on batch
                aff = []
                for index in range(nof_examples):
                    aff.append(np.sum(np.square(embeddings_array[index][:] - embeddings_array), 1))
                aff_binarized = binarize_affinity(aff, self.nof_images_per_age)
                # affinity matrix on

                # select triplets to train
                triplet = select_triplets(embeddings_array, self.nof_sampled_age, self.nof_images_per_age, self.delta)
                triplet_path_array = paths[triplet][:]
                triplet_label_array = labels[triplet][:]
                nof_triplets = len(triplet_path_array)
                print('%d triplets selected' % nof_triplets)

                # FIFO enqueue
                sess.run(self.enqueue_op, feed_dict={self.path_placeholder: triplet_path_array,
                                                     self.label_placeholder: triplet_label_array})
                # train on selected triplets
                nof_batches = int(np.ceil(nof_triplets * 3 / self.batch_size))
                for i in range(nof_batches):
                    batch_size = min(nof_triplets * 3 - i * self.batch_size, self.batch_size)
                    if self.step % 200 == 0:
                        _sum, loss, _ = sess.run(
                            [self.age_summary_op, self.age_loss, self.age_opt],
                            feed_dict={self.batch_size_placeholder: batch_size,
                                       self.age_affinity_watch: np.reshape(aff, [1, self.age_sampled_examples,
                                                                                self.age_sampled_examples,
                                                                                1]),
                                       self.age_affinity_watch_binarized: np.reshape(aff_binarized, [1,
                                                                                self.age_sampled_examples,
                                                                                self.age_sampled_examples,
                                                                                1]),
                                       self.nof_age_triplets_placeholder: nof_triplets,
                                       self.affinity_on_val: np.reshape(aff_val,
                                                                        [1, self.val_size, self.val_size, 1])})

                        # write in summary
                        summary_writer.add_summary(_sum, self.step)
                    else:
                        loss, _ = sess.run([self.age_loss, self.age_opt],
                                           feed_dict={self.batch_size_placeholder: batch_size})
                    progress(i + 1, nof_batches, str(triplet_selection) + 'th AGE Epoch' +str(inner_iter),
                             'Batches loss:' + str(loss))  # a command progress bar to watch training progress
                    self.step += 1
                    # save model
                    if self.step % 200000 == 0:
                        saver.save(sess, self.model_dir, global_step=self.step)

            # perform a validation on lfw
            val_paths = CACD.get_val(self.val_size)
            val_path_array = np.reshape(val_paths, (-1, 3))
            val_label_array = np.reshape(np.arange(self.val_size), (-1, 3))
            val_embeddings_array = np.zeros(shape=(self.val_size, self.embedding_size))

            # FIFO enqueue
            sess.run(self.enqueue_op,
                     feed_dict={self.path_placeholder: val_path_array,
                                self.label_placeholder: val_label_array})

            # forward propagation to get val embeddings
            print('Forward propagation on validation set')
            nof_batches = int(np.ceil(self.val_size / self.batch_size))
            for i in range(nof_batches):
                batch_size = min(self.val_size - i * self.batch_size, self.batch_size)
                emb, label = sess.run([self.id_embedding, self.label_batch],
                                      feed_dict={self.batch_size_placeholder: batch_size})
                val_embeddings_array[label, :] = emb
            aff_val = []
            for index in range(self.val_size):
                aff_val.append(np.sum(np.square(val_embeddings_array[index][:] - val_embeddings_array), 1))

def binarize_affinity(aff, k):
    temp = np.argsort(aff)
    ranks = np.arange(len(aff))[np.argsort(temp)]
    ranks[np.where(ranks > k)] = 255
    return ranks

def select_triplets(embeddings, nof_attr, nof_images_per_attr, delta):
    aff = []
    triplet = []
    for anchor_id in range(nof_attr * nof_images_per_attr):
        dist = np.sum(np.square(embeddings - embeddings[anchor_id]), 1)
        aff.append(dist)
        for pos_id in range(anchor_id + 1, (anchor_id // nof_images_per_attr + 1) * nof_images_per_attr):
            neg_dist = np.copy(dist)
            neg_dist[anchor_id:(anchor_id // nof_images_per_attr + 1) * nof_images_per_attr] = np.NAN
            neg_ids = np.where(neg_dist - dist[pos_id] < delta)[0]
            nof_neg_ids = len(neg_ids)
            if nof_neg_ids > 10:
                rand_id = np.random.randint(nof_neg_ids)
                neg_id = neg_ids[rand_id]
                triplet.append([anchor_id, pos_id, neg_id])
    np.random.shuffle(triplet)
    return triplet

if __name__ == '__main__':
    instance = DualReferenceFR()
    instance.train()


# cijfers ={1: 'Een', 2: 'Twee', 3: 'Drie', 4: 'Vier', 5: 'Vijf', 6: 'Zes', 7: 'Zeven', 8: 'Acht', 9: 'Negen', 10: 'Tien'}
#
# vraag1 = int(input('Voer een getal in tussen 1 en 10: ')[:2])
#
# while vraag1 > 10:
#     vraag1 = int(input('Nee, even serieus. Voer een getal in tussen 1 en 10: ')[:2])
# else:
#     print(cijfers[vraag1])
#
#
# print(vraag1, " is voor diegenen die niet weten hoe ze dat moeten spellen: ", cijfers[vraag1],'.', end='\n' "Is dat niet makkelijk?", sep='')

# namen = ['lies', 'jan', 'kees', 'mireille', 'koen', 'rob']
#
# for index, naam in enumerate(namen):
#     print('Nummer {} in de lijst is: {}'.format(index + 1, naam))
#
# #taak 4
# scores = {'lies': 6231, 'bas': 4322, 'kees': 6218, 'mireille':
# 4821, 'aniek': 6435, 'bert': 184, 'genius':10231}
# #print('Spelers:')
# # Methode 1: Alleen de sleutels worden afgedrukt
# for speler in scores:
#  print(speler)
# print()
# print('Scores:')
# # Methode 2: Alleen de waardes worden afgedrukt
# for score in scores.values():
#  print('Spelers met score')
#  print(score)
# print()
# # Methode 3: Zowel de sleutels als hun waardes worden afgedrukt
# for speler, score in scores.items():
#  print('{}: {}'.format(speler, score))

# scores = {'lies': 6231, 'bas': 4322, 'kees': 6218, 'mireille': 4821, 'aniek': 6435, 'bert': 184, 'genius': 10231}
#
# for speler, score in scores.items():
#     print('{:+^12}'.format(speler), '{:0>6}'.format(score))

# tuple_1 = (('Randy', 'Man', '25', 'Amsterdam')
# tuple_2 = ('Quincy', 'Man', '23', 'Amsterdam')
# tuple_3 = ('Maria', 'Vrouw', '55', 'Overrijsel')
# tuple_4 = ('Deborah', 'Vrouw', '32', 'Rotterdam')
# tuple_5 = ('Bob', 'Man', '14', 'Bunschoten-Spakenburg')
# tuple_6 = ('Laura', 'Vrouw', '28', 'Groningen'))

klas9A = (('Randy', 'Man', 25, 'Amsterdam'),\
          ('Quincy', 'Man', 23, 'Amsterdam'),\
          ('Maria', 'Vrouw', 55, 'Overrijsel'),\
          ('Deborah', 'Vrouw', 32, 'Rotterdam'),\
          ('Bob', 'Man', 14, 'Bunschoten-Spakenburg'),\
          ('Laura', 'Vrouw', 28, 'Groningen'))
for namen, geslacht, leeftijd, woonplaats in klas9A:
    if geslacht == 'Man':
        print('Zijn naam is {}, hij is {} leeftijd en woont in {}.'.format(namen, leeftijd, woonplaats))
    else:
        print('Haar naam is {}, zij is {} leeftijd en woont in {}.'.format(namen, leeftijd, woonplaats))


# experiment result wrapper

import numpy as np

import uuid
import h5py

from git import Repo
from pymongo import MongoClient

path_to_repo = '~/Documents/Thesis/latent_ssvm'
path_to_datafile = '/home/dmitry/Documents/Thesis/latent_ssvm/notebooks/experiment_data.hdf5'

class experiment(object):
    def __init__(self, f):
        self.f = f

    def __call__(self, description, *args, **kwargs):
        result = None
        try:
            result = self.f(*args, **kwargs)
        except:
            raise
        result.save(description)
        return result

class ExperimentResult(object):
    def __init__(self, data, meta, is_new=True):
        # meta information, comments, parameters
        # this will be stored in mongodb
        self.meta = meta
        if is_new:
            repo = Repo(path_to_repo)
            self.meta['commit_hash'] = repo.head.commit.hexsha
            # unique experiment identifier
            self.meta['id'] = uuid.uuid1().hex
        # generated data like scores per iteration, model parameters
        # stored in hdf5 file
        self.data = data

    def save_data(self):
        f = h5py.File(path_to_datafile, 'a', libver='latest')
        grp = f[self.meta['dataset_name']].create_group(self.meta['id'])
        for k in self.data.keys():
            grp.create_dataset(k, data=self.data[k])
        f.close()
        return grp.id.id

    def save_meta(self):
        client = MongoClient()
        client['lSSVM']['base'].insert(self.meta)
        client.disconnect()

    def save(self, description=''):
        self.description = description
        self.save_meta()
        self.save_data()
        return self.id

    @staticmethod
    def load(exp_id):
        client = MongoClient()
        meta = client['lSSVM']['base'].find_one({'id' : exp_id})
        f = h5py.File(path_to_datafile, 'r', libver='latest')
        grp = f[meta[u'dataset_name']][exp_id]
        data = {}
        for k in grp.keys():
            data[k] = np.empty(grp[k].shape)
            grp[k].read_direct(data[k])
        f.close()
        return ExperimentResult(data, meta, is_new=False)

    @property
    def description(self):
        return self.meta['description']

    @description.setter
    def description(self, description_):
        self.meta['description'] = description_

    @property
    def id(self):
        return self.meta['id']


import logging
import os

import requests
from google.cloud import datastore

logger = logging.getLogger(__name__)

def broadcast(check_result):
    notification = _create_notification_from_check_result(check_result)
    for token in _get_tokens():
        _send_notification(token, notification)

def _create_notification_from_check_result(check_result):
    title = 'In stock!' if check_result.in_stock else 'Out of stock...'
    return {
        'title': '[{}] {}'.format(check_result.store_key, title),
        'body': check_result.name,
        'click_action': check_result.url,
    }

def _get_tokens():
    client = datastore.Client()
    for token in client.query(kind='Token').fetch():
        yield token['token_id']

def _send_notification(to, notification):
    server_key = os.environ['X2SWITCH_FIREBASE_SERVER_KEY']
    headers = {
        'Authorization': 'key={}'.format(server_key),
    }
    message = {
        'to': to,
        'notification': notification,
    }
    response = requests.post('https://fcm.googleapis.com/fcm/send',
                             json=message, headers=headers)
    response.raise_for_status()
    logger.info('Notification result: %d %r',
                response.status_code, response.json())


import time, boto3, datetime, sys, csv, os, json

today = datetime.datetime.now().strftime("%Y%m%d")
year = datetime.datetime.now().strftime("%Y")
env = sys.argv[1] # input parameter

# open an AWS IAM boto session for the environment selected (live, test, security zone)
session = boto3.Session(profile_name=str(env))
iam_client = session.client("iam")

# open an AWS S3 boto session for the security tool instance
s3_client = boto3.client("s3")

# generate AWS IAM report (max 4 times per day)
try:
    iam_client.generate_credential_report()
except:
    pass

# wait for report generation
time.sleep(30)

# retrieve the report
report = iam_client.get_credential_report()

# output the results in a flat file
with open('/usr/local/src/security_scripts/outputs/'+today+'-aws-tmp'+ env +'-iam.csv', 'w+') as tmpFile:
    tmpFile.write(str(report['Content'],'utf-8'))

with open('/usr/local/src/security_scripts/outputs/' + today + '-aws-tmp' + env + '-iam.csv','r') as tmpFile, open('/usr/local/src/security_scripts/outputs/' + today + '-aws-' + env + '-iam.csv','w+') as outFile:
    lines = csv.reader(tmpFile)
    outFile.write("username,arn,user_creation_time,password_enabled,password_last_used,password_last_changed,password_next_rotation,mfa_active,access_key_1_active,access_key_1_last_rotated,access_key_1_last_used_date,access_key_1_last_used_region,access_key_1_last_used_service,access_key_2_active,access_key_2_last_rotated,access_key_2_last_used_date,access_key_2_last_used_region,access_key_2_last_used_service,cert_1_active,cert_1_last_rotated,cert_2_active,cert_2_last_rotated,groups\n")
    next(lines)
    for line in lines:
        # get all groups and policies attached for each user using AWS CLI as this is not provided into AWS credential report
        groups = ""
        #print(line[0])
        if line[0] != '<root_account>':
            attached_groups = iam_client.list_groups_for_user(UserName=line[0])['Groups']
            final_attached_groups = ''
            i = 0
            for group in attached_groups:
                if final_attached_groups != '' and attached_groups[i]['GroupName'] != '':
                    final_attached_groups = final_attached_groups + '|' + ''.join(attached_groups[i]['GroupName'])
                elif attached_groups[i]['GroupName'] != '':
                    final_attached_groups = ''.join(attached_groups[i]['GroupName'])
                i = i+1
            attached_managed_policies = iam_client.list_attached_user_policies(UserName = line[0])['AttachedPolicies']
            final_attached_managed_policies = ''
            i = 0
            for managed_policies in attached_managed_policies:
                if final_attached_managed_policies != '' and attached_managed_policies[i]['PolicyName'] != '':
                    final_attached_managed_policies = final_attached_managed_policies + '|' + ''.join(attached_managed_policies[i]['PolicyName'])
                elif attached_managed_policies[i]['PolicyName'] != '':
                    final_attached_managed_policies = ''.join(attached_managed_policies[i]['PolicyName'])
                i = i + 1
            attached_inline_policies = iam_client.list_user_policies(UserName = line[0])['PolicyNames']
            final_attached_inline_policies = ''.join(attached_inline_policies)
            groups = final_attached_groups + final_attached_inline_policies + final_attached_managed_policies
        #print(groups)
        outFile.write(line[0]+","+line[1]+","+line[2]+","+line[3]+","+line[4]+","+line[5]+","+line[6]+","+line[7]+","+line[8]+","+line[9]+","+line[10]+","+line[11]+","+line[12]+","+line[13]+","+line[14]+","+line[15]+","+line[16]+","+line[17]+","+line[18]+","+line[19]+","+line[20]+","+line[21]+","+groups+"\n")

# upload the result to s3
s3_client.upload_file('/usr/local/src/security_scripts/outputs/'+ today +'-aws-'+ env +'-iam.csv','securityzone-audit', 'aws-' + env + '-iam/' + year + '/' + today + '-aws-' + env + '-iam.csv')


from data_preprocess.data_filter import GetGuZhangData
import time
model_list=['cap_temp_1','cap_temp_2','cap_temp_3','convt_temp_1','convt_temp_2','convt_temp_3','motor_temp_1','motor_temp_2','motor_temp_3']
for model_kind in model_list:
    from model.xgboost.xgboostprocess import XgboostModel
    gz = GetGuZhangData('wfzc','A2',model_kind=model_kind)
    gz.process_data()
    xgbm = XgboostModel(jobname='xgb_model',fc='wfzc',fj='A2',model_kind=model_kind)
    xgbm.train_test_model()
    time.sleep(2)


"""新建一个组件."""
from typing import Dict, Any
from pmfp.const import (
    GOLBAL_PYTHON_VERSION
)
from pmfp.utils import (
    get_golang_version,
    get_node_version
)
from ._new_component import new_component
from ._new_doc import new_document
from ._new_env import new_env
from ._new_setup import new_setup
from ._new_readme import new_readme
from ._new_pb import new_pb
from ._new_es_script import new_es_script
from ._new_test import new_test

def new(config: Dict[str, Any], kwargs: Dict[str, Any]):
    """新建一个组件.

    Args:
        config (Dict[str, Any]): 项目配置
        kwargs (Dict[str, Any]): 组件配置

    Raises:
        e: 组件创建错误

    """
    if kwargs["language"] == "-":
        c_language = config["project-language"]
    else:
        c_language = kwargs["language"]

    if kwargs["to"] == "-":
        to = config["project-name"]
    else:
        to = kwargs["to"]
    c_name = kwargs["component_name"]
    if c_name in ("document", "doc"):
        new_document(config, c_language.lower())
    elif c_name == "env":
        new_env(config, c_language.lower())
    elif c_name == "readme":
        new_readme(config)
    elif c_name in ("setup", "cython_setup", "cython_numpy_setup", "cmd_setup"):
        new_setup(config, c_language, c_name)
    elif c_name in ("pb", "grpc", "grpc-streaming", "grpc-web-proxy"):
        if kwargs["rename"] == "-":
            rename = config["project-name"]
        elif kwargs["rename"] == "":
            rename = c_name
        else:
            rename = kwargs["rename"]
        new_pb(c_name=c_name, rename=rename, to=to, project_name=config["project-name"])
    elif c_name == "test":
        project_name = config["project-name"]
        if kwargs["rename"] == "-":
            rename = config["project-name"]
        elif kwargs["rename"] == "":
            rename = c_name
        else:
            rename = kwargs["rename"]
        new_test(c_language, project_name, rename)
    elif c_name == "es_script" and c_language == "Javascript":
        new_es_script(config)
    else:
        if not kwargs.get("kwargs"):
            kwargs["kwargs"] = {}
        if c_language == "Golang":
            kwargs["kwargs"]["language_version"] = get_golang_version() or "latest"
        elif c_language == "Python":
            kwargs["kwargs"]["language_version"] = GOLBAL_PYTHON_VERSION or "latest"
        elif c_language == "Javascript":
            kwargs["kwargs"]["language_version"] = get_node_version() or "latest"
        spl_name = c_name.split("-")
        c_category = spl_name[0]
        c_name = "".join(spl_name[1:])
        if kwargs["rename"] == "-":
            rename = config["project-name"]
        elif kwargs["rename"] == "":
            rename = c_name
        else:
            rename = kwargs["rename"]
        c_language = c_language.lower()
        c_path = f"{c_language}/{c_category}/{c_name}"
        test = kwargs["test"]
        try:
            nc_kwargs = kwargs.get("kwargs", {})
            print(f"创建组件{c_path},kwargs: {nc_kwargs}")
            new_component(
                config,
                path=c_path,
                to=to,
                rename=rename,
                test=test,
                **nc_kwargs)
        except Exception as e:
            print(f"组件{c_path}创建错误")
            raise e
        else:
            print(f"{c_path}创建为{rename}成功")


#6
print("Bienvenido al Sistema".center(50, "-"))

AUMENTO  = 0.08
DESCUENTO_POR_SERVICIO = 0.25

salario = 0

salario = int(input("ingrese su salario"))

salario = salario + int (salario * AUMENTO)
resta = (salario - DESCUENTO_POR_SERVICIO)
print("su salario es de:",salario)
print("su salario con descuento es de:",resta)


from window import main
import tkinter

class Furni:
    def __init__(self, name, x, y):
        self.name = name
        self.x = x
        self.y = y
        self.img = tkinter.PhotoImage(file='furniture/items/' + self.name)
        main.canvas.create_image(self.x, self.y, image=self.img)


# THIS FILE IS DEPRECATED; IT'S CONTENTS ARE NOW INCLUDED IN THE CORRESPONDING JUPYTER NOTEBOOK

# script for filtering pad chest dataset
#import numpy as np
import pandas as pd
import shutil

# paths
# path_external_drive = '/mnt/f/BIMCV-PadChest/unzipped/'
path_external_drive = '/mnt/c/Users/Jan/Daten/Geschäftlich/Capgemini/scripts/tmp/'
# path_pneumonia_files = 'relevant_files/pneumonia/'
path_pneumonia_files = ''
# csv_meta_data = 'remaining120000.csv'
# csv_meta_data = 'first40000.csv'
csv_meta_data = 'PADCHEST_chest_x_ray_images_labels_160K_01.02.19.csv'

# read meta data from csv to pandas dataframe
meta_data = pd.read_csv(path_external_drive + csv_meta_data, engine='python') # usecols= ['ImageID','ImageDir','Projection','Labels'], sep='(,{1})(?! )',

#print(meta_data.head)

meta_data['Labels'] = [label if label is not None and label is not [] and type(label) is not float else '' for label in meta_data['Labels']]
meta_data['Projection'] = [projection if projection is not None and projection is not [] else '' for projection in meta_data['Projection']]

# extract file names of relevant files
relevant_files = [True if 'normal' in data[31] and 'abnormal' not in data[31] and data[9] != 'L' else False for data in meta_data.values]

filtered_df = meta_data[relevant_files]
filtered_df.to_csv(path_external_drive + path_pneumonia_files + 'filtered_' + csv_meta_data)

number_of_files = len(filtered_df)
i = 0

print('found ' + str(number_of_files) + ' relevant files, starting to copy them')

# copying the files
for index, row in filtered_df.iterrows():
    i += 1
    # copy
    source = path_external_drive + str(row['ImageDir']) + '/' + str(row['ImageID'])
    destination = path_external_drive + path_pneumonia_files + str(row['ImageID'])
    shutil.copyfile(source, destination)
    print(str(i) + "/" + str(number_of_files) + " copied")

# --- TESTING ---
# meta_data2 = pd.read_csv('/mnt/c/Users/Jan/Daten/Geschäftlich/Capgemini/scripts/tmp/mini_csv_test_file.csv', usecols= ['ImageID','ImageDir','Projection','Labels'], sep='(,{1})(?! )', engine='python')
# meta_data2['Labels'] = [label if label is not None and not [] else '' for label in meta_data2['Labels']]
# relevant_files2 = [data for data in meta_data2 if 'fibrosis' in data]


"""
敏感词文本文件 filtered_words.txt，里面的内容为以下内容，当用户输入敏感词语时，则打印出 Freedom，否则打印出 Human Rights。
"""
"""
北京
程序员
公务员
领导
牛比
牛逼
你娘
你妈
love
sex
jiangge
"""
import sys

sensitive_words = ['北京', '程序员', '公务员', '领导', '牛逼', '牛比', '你娘', '你妈', 'love', 'sex', 'jiangge']

while (True):
    print("退出请输入:q")
    word = input("请输入检测词汇：")
    if word == ':q':
        sys.exit()
    if word in sensitive_words:
        print('Freedom')
    else:
        print('Human Rigjts')


r"""
===============================================================================
Submodule -- throat_diameter
===============================================================================

"""
from . import misc as _misc
import scipy as _sp

def weibull(geometry, shape, scale, loc, seeds=None, **kwargs):
    if seeds not in geometry:
        geometry['throat.seed'] = _sp.rand(geometry.Nt,)
    return _misc.weibull(geometry=geometry, shape=shape, scale=scale, loc=loc,
                         seeds=seeds)
weibull.__doc__ = _misc.weibull.__doc__

def normal(geometry, scale, loc, seeds=None, **kwargs):
    if seeds not in geometry:
        geometry['throat.seed'] = _sp.rand(geometry.Nt,)
    return _misc.normal(geometry=geometry, scale=scale, loc=loc,
                        seeds=seeds)
normal.__doc__ = _misc.normal.__doc__

def generic(geometry, func, seeds=None, **kwargs):
    if seeds not in geometry:
        geometry['throat.seed'] = _sp.rand(geometry.Nt,)
    return _misc.generic(geometry=geometry, func=func, seeds=seeds)
generic.__doc__ = _misc.generic.__doc__

def random(geometry, seed=None, num_range=[0, 1], **kwargs):
    r"""
    Assign throat sizes from a random

    Parameters
    ----------
    geometry : OpenPNM Geometry object
        The Geometry object to which this model will apply.  This is necessary
        to determine the length of the array to generate.

    seed : int
        The starting seed value to send to Scipy's random number generator.
        The default is None, which means different distribution is returned
        each time the model is run.

    num_range : list
        A two element list indicating the low and high end of the returned
        numbers.  The default is [0, 1] but this can be adjusted to produce
        throat sizes directly; for instance pores between 10 and 100 um can be
        generated with ``num_range = [0.00001, 0.0001]``.
    """
    return _misc.random(N=geometry.Nt, seed=seed, num_range=num_range)

def cylinder(geometry, tsd_name, tsd_shape, tsd_loc, tsd_scale,
             throat_seed='throat.seed', tsd_offset=0, **kwargs):
    r"""
    Calculate throat diameter from seeds for a cylindrical throat

    Parameters
    ----------
    geometry : OpenPNM Geometry Object
        The Geometry object which this model is associated with. This controls
        the length of the calculated array, and also provides access to other
        necessary geometric properties.

    tsd_name : string
        The name of the statistical distribution to use. This model uses the
        Scipy.stats module, so any of the distributions available there are
        suitable options.

    tsd_shape, loc and scale : scalars
        The parameters to send to the selected statistics model.  Most of the
        Scipy.stats models accept the same keyword arguments.  Note that the
        psd_ prefix is added by OpenPNM to indicate 'pore size distribution'.

    tsd_offset : scalar
        Controls the minimum value in the pore size distribution by shifting
        the entire set of values by the given offset.  This is useful for
        avoiding pore sizes too close to zero.

    Notes
    -----
    This pore-scale model is deprecated.  Use ``weibull``, ``normal`` or
    ``generic`` to get produce pore sizes distributions.
    """
    import scipy.stats as spst
    prob_fn = getattr(spst, tsd_name)
    P = prob_fn(tsd_shape, loc=tsd_loc, scale=tsd_scale)
    value = P.ppf(geometry[throat_seed]) + tsd_offset
    return value

def equivalent_circle(geometry, throat_area='throat.area', **kwargs):
    r"""
    Calculates the diameter of a cirlce with same area as the throat.

    Parameters
    ----------
    geometry : OpenPNM Geometry object
        The Geometry object which contains the throat area values

    thorat_area : string
        The dictionary key to the throat area values
    """
    areas = geometry[throat_area]
    value = 2*_sp.sqrt(areas/_sp.pi)  # 64 bit sqrt doesn't work!
    return value

def minpore(network, geometry, factor=1, **kwargs):
    r"""
    Assign the throat diameter to be equal to the smallest connecting pore
    diameter. If zero (in case of boundaries) take it to be the maximum of
    the connecting pore diameters

    Parameters
    ----------
    factor : float < 1
        A factor between 0 and 1 to further constrict the throat size
        calculcated by the function.

    """
    gTs = geometry.throats()
    nTs = geometry.map_throats(network, gTs)
    pDs = network["pore.diameter"][network["throat.conns"][nTs]]
    value = _sp.amin(pDs, axis=1)*factor
    value[value == 0.0] = _sp.amax(pDs, axis=1)[value == 0.0]
    return value


#!/usr/bin/env python
import os, json, argparse, setlog, parsegenomes, mash, shutil, re
import copyseqsql, makeseqsql, glob, seqsql2fa, subprocess
import multiprocessing as mp
import makehmmsql, getgenematrix, getgenes, concatmsa, ete3helper
import numpy as np
#from scipy.cluster.vq import kmeans2

###
##### Workflow 1: Build query + reference phylogeny from scratch
###

def catTrees(flist,outfile):
    with open(outfile,"w") as ofil:
        for fname in flist:
            if os.path.exists(fname):
                with open(fname,"r") as ifil:
                    for line in ifil:
                        ofil.write(line)
            else:
                log.warning("Tree file not found for MLST: %s"%fname)
    return outfile

def runmafft(finput,output,thread=1,maxit=300,localpair=False,options="",rename="",fast=False):
    if localpair:
        options += " --localpair"
    if rename:
        #Rename fasta headers by ignoring everything after "rename" character
        with open(finput,"r") as infil, open(finput+".renamed","w") as ofil:
            for line in infil:
                if line.startswith(">") and rename in line:
                    ofil.write(line.strip().split(rename)[0]+"\n")
                else:
                    ofil.write(line)
        finput = finput+".renamed"
    if fast:
        cmd = "mafft --quiet --thread %d --treeout %s %s"%(thread,options,finput)
    else:
        cmd = "mafft --quiet --thread %d --localpair --treeout --maxiterate %d %s %s"%(thread,maxit,options,finput)
    cmd = cmd.split()
    with open(os.devnull,"w") as devnull, open(output,"w") as outfil:
        try:
            subprocess.call(cmd, stdout=outfil, stderr=devnull)
            log.info("Finished alignment %s"%output)
        except Exception as e:
            log.error("Failed to run mafft %s (%s)"%(finput,e))

def runtrimal(finput,output,method="automated1"):
    cmd = "trimal -in %s -out %s -%s"%(finput,output,method)
    cmd = cmd.split()
    with open(os.devnull,"w") as devnull:
        try:
            subprocess.call(cmd, stdout=devnull, stderr=devnull)
            log.info("Finished trimming %s"%output)
        except Exception as e:
            log.error("Failed to run trimall %s (%s)"%(finput,e))

def processmlst(indir,aligndir,cpu=1,parallel=True,trim=False,fast=False):
    flist = glob.glob(os.path.join(indir,"*.fna"))
    if not os.path.exists(aligndir):
        os.makedirs(aligndir)

    if cpu > 1 and parallel:
        pool = mp.Pool(cpu)
        try:
            for filename in flist:
                fname = os.path.split(filename)[-1]
                if trim:
                    pool.apply_async(runtrimal, args=(filename, os.path.join(aligndir,fname)))
                else:
                    pool.apply_async(runmafft, (filename, os.path.join(aligndir,fname)),dict(fast=fast))
        finally:
            pool.close()
            pool.join()
    else:
        for filename in flist:
            fname = os.path.split(filename)[-1]
            if trim:
                runtrimal(filename, os.path.join(aligndir,fname))
            else:
                runmafft(filename, os.path.join(aligndir,fname),thread=cpu,fast=fast)

def hmmsearch(outname,hmmdb,infile,mcpu=1,cut="tc"):
    log.info("Searching sequences against hmmdb: %s"%hmmdb)
    cmd=["hmmsearch", "--domtblout", outname, "--noali", "--notextw", hmmdb, infile]
    if mcpu>1:
        cmd[1:1] = ["--cpu", str(mcpu)]
    if cut and cut.lower() in ("ga","tc","nc"):
        cmd[1:1] = ["--cut_%s"%cut.lower()]
    with open(os.devnull,"w") as devnull:
        subprocess.call(cmd, stdout=devnull, stderr=devnull)

def buildseqdb(selorgs,queryseqs,resultdir,sourcedb):
    #Get list of all selected and outgroups
    reflist = list(selorgs.get("selspecies",[]))
    reflist.extend(selorgs.get("seloutgroups",[]))

    #Copy from reference database to query database
    seqsdb = os.path.join(resultdir,"seqs.db")
    copyseqsql.copydb(queryseqs,sourcedb,seqsdb)

def getorgs(resultdir,mashresult,skip="",IGlimit=50,OGlimit=1,minorgs=25):
    #Get user selection if exists
    usersel = os.path.join(resultdir,"userlist.json")
    autosel = os.path.join(resultdir,"autoOrglist.json")
    selection = False
    if os.path.exists(usersel):
        with open(usersel,"r") as fil:
            selection = json.load(fil)
    elif os.path.exists(autosel):
        with open(autosel,"r") as fil:
            selection = json.load(fil)
    elif "skip2" in skip:
        # commonrank = mashresult.get("commonrank",["phyl","",""])
        # if commonrank[0]:
        #     selection = {"selspecies":[x.get("id","") for x in mashresult.get("reforgs",[])[:IGlimit] if str(x.get(commonrank[0]+"id","")) in str(commonrank[1])],
        #              "seloutgroups":[x.get("id","") for x in mashresult.get("outgroups",[])[:OGlimit]]}
        # else:
        qorgs = mashresult.get("queryorgs",[])
        selrefs = mashresult.get("reforgs",[])[:IGlimit-len(qorgs)-OGlimit]
        selection = {"selspecies":[x.get("id","") for x in selrefs]}
        orgrecs = {i:v for i,v in enumerate(qorgs+selrefs)}
        # log.debug("Issue_orgrecs: %s"%orgrecs)
        commonrank = mash.getcommongroup(orgrecs)
        # log.debug("Issue_commonrank: %s"%commonrank)
        OGorgs = mash.getoutgrouporgs(commonrank,mashresult.get("reforgs",[]),glimit=OGlimit)
        # log.debug("Issue_OGorgs: %s"%OGorgs)
        selection["seloutgroups"] = [x.get("id","") for x in OGorgs[:OGlimit]]
        # if not len(selection["selspecies"]) > minorgs:
        #     selection = {"selspecies":[x.get("id","") for x in mashresult.get("reforgs",[])[:minorgs]],
        #              "seloutgroups":[x.get("id","") for x in mashresult.get("outgroups",[])[:OGlimit]]}

        with open(autosel,"w") as fil:
            json.dump(selection,fil)

    return selection

def getmlstselection(resultdir,mlstpriority,maxmlst=100,minmlst=10,skip="",ignoreorgs=None,concat=True):
    if not ignoreorgs:
        ignoreorgs = []
    usersel = os.path.join(resultdir,"usergenes.json")
    autosel = os.path.join(resultdir,"autoMlstlist.json")
    selection = False
    delorgs = set()
    if os.path.exists(usersel):
        with open(usersel,"r") as fil:
            temp = json.load(fil)
            selection = temp["selection"]
            delorgs = temp["delorgs"]
            concat = temp.get("mode","concatenated")
            if concat != "concatenated":
                concat = False
            return selection, list(delorgs), concat

    if os.path.exists(autosel):
        with open(autosel,"r") as fil:
            temp = json.load(fil)
            selection = temp["selection"]
            delorgs = temp["delorgs"]
    else:
        minrecs = [x for x in mlstpriority if not len(x["orgdel"])]
        log.info("MLST genes in all orgs = %s"%len(minrecs))
        if len(minrecs) >= minmlst:
            recs = minrecs[:maxmlst]
        else:
            recs = [x for x in mlstpriority if not any([org in x["orgdel"] for org in ignoreorgs])][:minmlst]
        selection = [x["acc"] for x in recs]
        for x in recs:
            delorgs = delorgs | set(x["orgdel"])
        with open(autosel,"w") as fil:
            json.dump({"selection":selection,"delorgs":list(delorgs)},fil)

    return selection, list(delorgs), concat

def runIQtree(outdir,infasta,partfile="",cpu=1,model="GTR",bs=0,fout="speciestree",titlesep="",outgroup=""):
    #If title seperator exists rename titles in file first
    if titlesep:
        tf = infasta+".renamed"
        with open(tf,"w") as ofil, open(infasta,"r") as ifil:
            for line in ifil:
                lineout = line.strip()
                if line.startswith(">") and titlesep in line:
                    lineout = line.strip().split(titlesep)[0]
                ofil.write(lineout+"\n")
        infasta = tf
    log.info("Building tree: model=%s, bootstrap=%s, file=%s"%(model,bs,fout))
    cmd=["iqtree", "-quiet", "-nt", str(cpu), "-s", infasta, "-m", model,"-pre",os.path.join(outdir,fout)]
    if bs:
        cmd.extend(["-bb",str(bs)])
    if outgroup:
        #get full title of first outgroup to root tree
        with open(infasta,"r") as ifil:
            for line in ifil:
                if line.startswith(">") and outgroup in line:
                    cmd.extend(["-o", "%s"%line[1:].strip()])
                    break
    if partfile:
        cmd.extend(["-spp", partfile])
    with open(os.devnull,"w") as devnull:
        subprocess.call(cmd,stdout=devnull,stderr=devnull)

def runAstral(resultdir,intree,outtree,astjar=""):
    log.info("Started running Astral: %s"%(intree))
    if not astjar or not os.path.realpath(astjar):
        astjar = os.path.join(os.path.dirname(__file__),"astral.5.5.9.jar")
    cmd=["java", "-Xmx3000M", "-jar", str(astjar), "-i", intree, "-o", outtree]
    with open(os.path.join(resultdir,"astral.log"),"w") as astlog:
        subprocess.call(cmd,stdout=astlog,stderr=astlog)

def colphylogeny(resultdir,trimdir,cpu=1,model="MFP",bs=0):
    treedir = os.path.join(resultdir,"trees")
    if not os.path.exists(treedir):
        os.makedirs(treedir)
    flist = glob.glob(os.path.join(trimdir,"*.fna"))
    #build all trees
    if cpu > 1:
        pool = mp.Pool(cpu)
        try:
            for filename in flist:
                pool.apply_async(runIQtree, (treedir, filename), dict(model=model,bs=bs,fout=os.path.split(filename)[-1],titlesep="|"))
        finally:
            pool.close()
            pool.join()
    else:
        for filename in flist:
            runIQtree(treedir, filename, fout=os.path.split(filename)[-1], bs=bs, model=model, titlesep="|")

    #Combine all trees using ASTRAL
    flist = glob.glob(os.path.join(treedir,"*.treefile"))
    alltrees = catTrees(flist,os.path.join(treedir,"alltrees.tree"))
    coltree = os.path.join(treedir,"summaryTree.tree")
    runAstral(resultdir,alltrees,coltree)

def concatphylogeny(resultdir,concatfasta,partfile,cpu=1,model="GTR",bs=0,outgroup=""):
    trimdir = os.path.join(resultdir,"mlst_trimmed")
    log.info("Concatenating all MLST genes")
    concatmsa.concatmsa(os.path.join(trimdir,"*.fna"),concatfasta,partfile,fsplit=" ",checktype=False)
    #Build tree from concatenated fasta
    treedir = os.path.join(resultdir,"trees")
    if not os.path.exists(treedir):
        os.makedirs(treedir)
    runIQtree(treedir,concatfasta,partfile=partfile,cpu=cpu,fout="concatTree.tree",bs=bs,model=model)

def screenmlst(mlstdir,aligndir,cpu=1,mingenes=50):
    """Additional screen for outlier genes with highly discordant parsimony trees"""
    guidetrees = glob.glob(os.path.join(mlstdir,"*.tree"))
    #combine all to tree file
    hkgenes = []
    alltrees = os.path.join(aligndir,"guidetrees.tree")
    with open(alltrees,"w") as gtfil:
        for fname in guidetrees:
            with open(fname,"r") as tfil:
                hk = os.path.split(fname)[-1]
                hk = hk[:hk.find(".")]
                hkgenes.append(hk)
                gtfil.write("".join([x.strip() for x in tfil]).replace(";","")+";\n")
    #Run robinson foulds distance calculation
    cmd = ["iqtree","-nt",str(cpu),"-rf_all",alltrees]
    with open(os.devnull,"w") as devnull:
        subprocess.call(cmd,stdout=devnull,stderr=devnull)

    #Read distance matrix and store median values for each gene
    treedists = []
    with open(alltrees+".rfdist","r") as fil:
        line = fil.next()
        for line in fil:
            temp = np.median([int(x) for x in line.strip().split()[1:] if int(x)>0])
            treedists.append(temp)

    #segment tress into 2 partitions using k-means
    #km,kmgroup = kmeans2(np.array(treedists),2)
    #Ensure group with highest distance is marked as group 1, group 0 is highest group
    #if km[0] > km[1]:
    #    kmgroup = np.ones(len(kmgroup))-kmgroup
    #inds = [i for i,x in enumerate(kmgroup) if x>0]
    #get hkgenes with highest group designation
    # return kmgroup,hkgenes,treedists

    #Assume normal distribution and highlight trees outside 1 standard deviation (~ 15% of distant genes)
    threshold = np.mean(treedists)+np.std(treedists)
    inds = np.array([int(x>=threshold) for x in treedists])

    hkhigh = sorted(np.array(hkgenes)[inds>0], key=lambda x: treedists[hkgenes.index(x)])

    #if lowest group is under min genes remove genes from highest group
    mindiff = mingenes - (len(hkgenes)-len(hkhigh))
    if mindiff > 0:
        hkhigh = hkhigh[mindiff+1:]

    #Mask screened MLST genes
    for hk in hkhigh:
        fname = os.path.join(aligndir,hk+".fna")
        os.rename(fname,fname+".removed")

    #write table of genes and median distances
    fname = os.path.join(aligndir,"removedsumamry.tsv")
    with open(fname,"w") as fil:
        for i,hk in enumerate(hkgenes):
            fil.write("%s\t%s\t%s\n"%(hk,treedists[i],inds[i]))

    return hkhigh

def startwf1(indir,resultdir,checkpoint=False,concat=False,mashmxdist=0.5,cpu=1,
             skip="",refdb="",hmmdb="",rnadb="",maxmlst=100,model="GTR",bs=0,
             kf=False,maxorg=50,filtMLST=True,fast=False,minmlst=10):
    """WORKFLOW 1: Build phylogeny from scratch"""
    if not checkpoint:
        checkpoint = "w1-0"
    queryseqs = os.path.join(resultdir,"queryseqs")

    #Parse all inputs
    if checkpoint == "w1-0":
        log.info("JOB_STATUS::Parsing all genomes...")
        if parsegenomes.parseall(indir,queryseqs):
            checkpoint = "w1-1"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::5/100")
        else:
            log.error("Problem parsing input genomes")
            return False

    #Run MASH distances
    mashresult = False
    if checkpoint == "w1-1":
        log.info("JOB_STATUS::Running MASH ANI estimation against reference sequences...")
        mashresult = mash.getdistances(queryseqs,resultdir,cpu=cpu,maxdist=mashmxdist)
        if mashresult:
            checkpoint = "w1-2"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::10/100")
        else:
            log.error("MASH distance failed")
            return False

    #Get set of organisms to build seq DB
    selorgs = False
    if checkpoint == "w1-2":
        log.info("JOB_STATUS::Loading mash results...")
        if mashresult:
            pass
        elif os.path.exists(os.path.join(resultdir,"reflist.json")):
            with open(os.path.join(resultdir,"reflist.json"),"r") as fil:
                mashresult = json.load(fil)
                log.info("Loading mash results...")
        else:
            log.error("No Mash results to process")
            return False

        selorgs = getorgs(resultdir,mashresult,skip=skip,IGlimit=maxorg,minorgs=25)
        if selorgs:
            checkpoint = "w1-3"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::15/100")
        else:
            checkpoint = "w1-STEP2"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)

    if checkpoint == "w1-STEP2":
        log.info("JOB_STATUS:: Waiting for selected organisms")
        return "waiting"

    #Copy reference sequence database and add query organisms
    orgdb = os.path.join(resultdir,"refquery.db")
    # selorgs = {}
    if not selorgs and os.path.exists(os.path.join(resultdir,"userlist.json")):
        with open(os.path.join(resultdir,"userlist.json"),"r") as fil:
            selorgs = json.load(fil)
    elif not selorgs and os.path.exists(os.path.join(resultdir,"autoOrglist.json")):
        with open(os.path.join(resultdir,"autoOrglist.json"),"r") as fil:
            selorgs = json.load(fil)
    if checkpoint == "w1-3":
        log.info("JOB_STATUS:: Collating selected genomes...")
        #Clear old db if exists
        if os.path.exists(orgdb):
            os.remove(orgdb)

        flist = list(selorgs["selspecies"])
        flist.extend(selorgs["seloutgroups"])
        flist = [x for x in flist if "query" not in x.lower()]
        log.info("refdb: %s"%refdb)
        if copyseqsql.copydb(flist,refdb,orgdb):
            #get file list of query orgs and add to db
            seqlist = glob.glob(os.path.join(queryseqs,"*.fna"))
            allorgs = makeseqsql.runlist(seqlist,orgdb,True,"",True)
            if allorgs:
                checkpoint = "w1-4"
                log.info("JOB_CHECKPOINT::%s"%checkpoint)
        else:
            log.error("Failed at copydb")

    #Run HMM searches on query organism
    if checkpoint == "w1-4":
        if not orgdb:
            log.error("No querydb rerun at checkpoint w1-3")
            return False
        #Write newly added seqences to file
        naseqs = os.path.join(resultdir,"addedseqs.fna")
        aaseqs = os.path.join(resultdir,"addedseqs.faa")
        seqsql2fa.writefasta(orgdb,aaseqs,False,"",True)
        seqsql2fa.writefasta(orgdb,naseqs,True,"",True)

        #Run HMM searches
        log.info("JOB_STATUS:: Searching for MLST genes in query sequences...")
        log.info("JOB_PROGRESS::25/100")
        if not hmmdb:
            hmmdb = os.path.join(os.path.dirname(os.path.realpath(__file__)),"reducedcore.hmm")
        if not rnadb:
            rnadb = os.path.join(os.path.dirname(os.path.realpath(__file__)),"barnap_bact_rRna.hmm")
        hmmsearch(aaseqs+".domhr",hmmdb,aaseqs,mcpu=cpu)
        hmmsearch(naseqs+".domhr",rnadb,naseqs,mcpu=cpu)

        #Add HMMresults to refdb and cleanup unused seqs
        log.info("Adding query HMM results to database")
        status = makehmmsql.run(aaseqs+".domhr",orgdb)
        if not status:
            #If no genes were found report error
            log.error("No MLST genes could be found. Stop processing")
        log.info("Adding query RNA results to database")
        status = makehmmsql.run(naseqs+".domhr",orgdb,rna=True)
        os.remove(aaseqs)
        os.remove(naseqs)
        checkpoint = "w1-5"
        log.info("JOB_CHECKPOINT::%s"%checkpoint)

    #Calculate mlst core genes and export
    mlstdir = os.path.join(resultdir,"mlstgenes")
    mlstpriority = os.path.join(resultdir,"mlstpriority.json")
    genematjson = os.path.join(resultdir,"mlstmatrix.json")
    dndsfile = os.path.join(os.path.dirname(os.path.realpath(__file__)),"dnds.json")
    mlstselection = []
    delorgs = []
    if checkpoint == "w1-5":
        allquery = [os.path.splitext(os.path.split(x)[-1])[0].replace(" ","_") for x in glob.glob(os.path.join(resultdir,"queryseqs","*.fna"))]
        ignoreorgs = list(selorgs.get("seloutgroups",[]))
        ignoreorgs.extend(allquery)

        #getmlstgenes.findsingles(orgdb,maxgenes=500,outdir=mlstdir)
        if os.path.exists(mlstpriority) and os.path.exists(genematjson):
            with open(genematjson,"r") as mfil, open(mlstpriority,"r") as pfil:
                mlstpriority = json.load(pfil)
                temp = json.load(mfil)
                genemat = temp["counts"]
                orgs = temp["orgs"]
                del temp
        else:
            genemat,orgs,mlstpriority = getgenematrix.getmat(orgdb,pct=0.5,pct2=1.0,bh=True,rna=True,savefil=genematjson,prifile=mlstpriority,dndsfile=dndsfile,ignoreorgs=allquery)
        mlstselection, delorgs, concat = getmlstselection(resultdir,mlstpriority,maxmlst,ignoreorgs=ignoreorgs,concat=concat,minmlst=minmlst)
        if "skip3" in skip.lower() or os.path.exists(os.path.join(resultdir,"usergenes.json")):
            #Export selected genes to mlst folder
            log.info("JOB_STATUS:: Writing MLST genes...")
            getgenes.writeallgenes(orgdb,mlstselection,delorgs,outdir=mlstdir,outgroups=selorgs.get("seloutgroups",None),pct=0.5,rename=True)
            checkpoint = "w1-6"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
        else:
            checkpoint = "w1-STEP3"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)

    if checkpoint == "w1-STEP3":
        log.info("JOB_STATUS:: Waiting for selected MLST genes")
        return "waiting"

    ## Align and trim all MLST genes
    aligndir = os.path.join(resultdir,"mlst_aligned")
    trimdir = os.path.join(resultdir,"mlst_trimmed")

    if checkpoint == "w1-6":
        log.info("JOB_STATUS:: Aligning MLST genes")
        log.info("JOB_PROGRESS::30/100")
        #align all
        processmlst(mlstdir,aligndir,cpu=cpu,fast=fast)
        checkpoint = "w1-6b"
        log.info("JOB_CHECKPOINT::%s"%checkpoint)

    if checkpoint == "w1-6b":
        if filtMLST and filtMLST != "False":
            #Extra screen of MLST genes to remove outliers based on starting tree distance
            log.info("JOB_STATUS:: Screening for inconsistent MLST genes")
            log.info("JOB_PROGRESS::55/100")
            excludemlst = screenmlst(mlstdir,aligndir,cpu=cpu)
            log.info("JOB_STATUS:: Excluded MLST genes: %s"%excludemlst)

        #trim all
        log.info("JOB_STATUS:: Trimming alignments")
        log.info("JOB_PROGRESS::65/100")
        processmlst(aligndir,trimdir,cpu=cpu,trim=True)

        checkpoint = "w1-7"
        log.info("JOB_CHECKPOINT::%s"%checkpoint)

    treedir = os.path.join(resultdir,"trees")
    finishedtree = ""
    #Build trees
    if checkpoint == "w1-7":
        log.info("JOB_PROGRESS::75/100")
        if concat and concat != "False":
            log.info("JOB_STATUS:: Running concatenated supermatrix phylogeny")
            concatfasta = os.path.join(resultdir,"concatMLST.fasta")
            partfile = os.path.join(resultdir,"nucpartition.txt")
            concatphylogeny(resultdir, concatfasta, partfile,cpu=cpu,model=model,bs=bs)
            checkpoint = "w1-F"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
        else:
            log.info("JOB_STATUS:: Running coalescent tree phylogeny")
            colphylogeny(resultdir,trimdir,cpu=cpu,model=model,bs=bs)
            checkpoint = "w1-F"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)

    if checkpoint == "w1-F":
        finaltree = os.path.join(resultdir,"final.tree")
        if os.path.exists(os.path.join(treedir,"concatTree.tree.treefile")) and concat:
            finishedtree = os.path.join(treedir,"concatTree.tree.treefile")
        elif os.path.exists(os.path.join(treedir,"summaryTree.tree")):
            finishedtree = os.path.join(treedir,"summaryTree.tree")
        #Copy final tree to root dir
        if finishedtree:
            fmat = 2 if bs or not concat else 5
            log.debug("Saving final tree... %s"%fmat)
            if not ete3helper.rerootTree(finishedtree,finaltree,fmat=fmat):
                shutil.copy(finishedtree,finaltree)
        else:
            log.error("Could not find final tree. Tree building failed")

        #Job finished? do some cleanup
        if not kf:
            if os.path.exists(orgdb):
                os.remove(orgdb)
            shutil.rmtree(os.path.join(resultdir,"queryseqs"))
            shutil.rmtree(os.path.join(resultdir,"mlst_trimmed"))
            shutil.rmtree(os.path.join(resultdir,"mlstgenes"))
            for oldfil in glob.glob(os.path.join(treedir,"*.model")):
                os.remove(oldfil)

def raxmlEPA(outdir,inalgn,reftree):
    #Run raxml EPA
    cmd="raxmlHPC-SSE3 -f v -m GTRGAMMA -p 12345 -w %s -t %s -s %s -n %s"%(outdir,reftree,inalgn,os.path.split(reftree)[-1])
    log.debug("Starting: %s"%cmd)
    try:
        with open(os.devnull,"w") as devnull:
            subprocess.call(cmd.split(),stdout=devnull)
        log.debug("RAxML-EPA: finished %s"%inalgn)
        return True
    except subprocess.CalledProcessError as e:
        log.error("RAxML-EPA: error, could not process %s - %s"%(inalgn,e))
        return False

def addalltrees(inlist,famrefdir,outdir,cpu=1):
    if cpu==1:
        for x in inlist:
            fname = os.path.split(x)[1]
            ftree = os.path.join(famrefdir,fname+".tree")
            raxmlEPA(outdir,x,ftree)
    else:
        pool = mp.Pool(cpu)
        try:
            for x in inlist:
                fname = os.path.split(x)[1]
                ftree = os.path.join(famrefdir,fname+".tree")
                pool.apply_async(raxmlEPA, (outdir,x,ftree))
        finally:
            pool.close()
            pool.join()

def addallalign(inlist,famrefdir,outdir,cpu=1,fast=False):
    if cpu==1:
        for x in inlist:
            fname = os.path.split(x)[1]
            addopt = "--add %s"%x
            finput = os.path.join(famrefdir,fname)
            output = os.path.join(outdir,fname)
            runmafft(finput,output,options=addopt,fast=fast)
    else:
        pool = mp.Pool(cpu)
        try:
            for x in inlist:
                fname = os.path.split(x)[1]
                addopt = "--add %s"%x
                finput = os.path.join(famrefdir,fname)
                output = os.path.join(outdir,fname)
                pool.apply_async(runmafft, (finput, output),dict(options=addopt,fast=fast))
        finally:
            pool.close()
            pool.join()

def getreference(mashresult,refdir):
    reffams = [str(x).lower() for x in os.listdir(refdir)]

    fams = list(set([q.get("familyname","N/A") for q in mashresult["queryorgs"]]))
    qfam = str(fams[0]).lower()
    if len(fams) == 1 and qfam in reffams:
        #get exact case as folder
        family = [fam for fam in os.listdir(refdir) if qfam == fam.lower()]
        return str(fams[0])
    elif len(fams) > 1:
        famlist = [(x,mashresult["queryorgs"][i].get("familyname","N/A")) for i,x in enumerate(mashresult["orglist"])]
        log.error("Multiple families detected in query geneomes: %s" % "\t".join( ["%s = %s |"%x for x in famlist]) )
    else:
        log.error("No matching family detected")

def startwf2(indir,resultdir,refdir="",checkpoint=False,reference="",model="GTR",bs=0,kf=False,maxmlst=100,
             cpu=1,mashmxdist=0.5,rnadb="",hmmdb="",fast=False):
    """WORKFLOW 2: Get all query genomes and identify reference tree to add sequences to"""
    if not checkpoint:
        checkpoint = "w2-0"
    if reference:
        checkpoint = "w2-3"
    queryseqs = os.path.join(resultdir,"queryseqs")

    #Parse all inputs
    if checkpoint == "w2-0":
        log.info("JOB_STATUS::Parsing all genomes...")
        if parsegenomes.parseall(indir,queryseqs):
            checkpoint = "w2-1"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::5/100")
        else:
            log.error("Problem parsing input genomes")
            return False

    #Run MASH distances
    mashresult = False
    if checkpoint == "w2-1":
        log.info("JOB_STATUS::Running MASH ANI estimation against reference sequences...")
        mashresult = mash.getdistances(queryseqs,resultdir,cpu=cpu,maxdist=mashmxdist)
        if mashresult:
            checkpoint = "w2-2"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::10/100")
        else:
            log.error("MASH distance failed")
            return False

    #Detect reference group
    if checkpoint == "w2-2":
        log.info("JOB_STATUS::Loading mash results...")
        if mashresult:
            pass
        elif os.path.exists(os.path.join(resultdir,"reflist.json")):
            with open(os.path.join(resultdir,"reflist.json"),"r") as fil:
                mashresult = json.load(fil)
                log.info("Loading mash results...")
        else:
            log.error("No Mash results to process")
            return False

        reference = getreference(mashresult,refdir)
        if reference:
            checkpoint = "w2-3"
            log.info("JOB_STATUS::Detected reference = %s"%reference)
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.info("JOB_PROGRESS::15/100")
        else:
            checkpoint = "w2-F"
            log.info("JOB_CHECKPOINT::%s"%checkpoint)
            log.error("No matching reference found for all query organisms")

    orgdb = os.path.join(resultdir,"allseqs.db")
    seqlist = glob.glob(os.path.join(queryseqs,"*.fna"))
    #Add to db and run HMM searches. Use model in reference folder, or use global models
    if checkpoint == "w2-3":
        #get file list of query orgs and add to db
        allorgs = makeseqsql.runlist(seqlist,orgdb,True,"",False)

        #Write newly added seqences to file
        naseqs = os.path.join(resultdir,"addedseqs.fna")
        aaseqs = os.path.join(resultdir,"addedseqs.faa")
        seqsql2fa.writefasta(orgdb,aaseqs,False,"",True)
        seqsql2fa.writefasta(orgdb,naseqs,True,"",True)

        #Run HMM searches
        log.info("JOB_STATUS:: Searching for MLST genes in query sequences...")
        log.info("JOB_PROGRESS::25/100")

        if not hmmdb:
            if os.path.exists(os.path.join(refdir,reference,"core.hmm")):
                hmmdb = os.path.join(refdir,reference,"core.hmm")
            else:
                hmmdb = os.path.join(os.path.dirname(os.path.realpath(__file__)),"reducedcore.hmm")
        if not rnadb:
            rnadb = os.path.join(os.path.dirname(os.path.realpath(__file__)),"barnap_bact_rRna.hmm")

        hmmsearch(aaseqs+".domhr",hmmdb,aaseqs,mcpu=cpu)
        hmmsearch(naseqs+".domhr",rnadb,naseqs,mcpu=cpu)

        #Add HMMresults
        log.info("Adding query HMM results to database")
        status = makehmmsql.run(aaseqs+".domhr",orgdb)
        if not status:
            #If no genes were found report error
            log.error("No MLST genes could be found. Stop processing")
        log.info("Adding query RNA results to database")
        status = makehmmsql.run(naseqs+".domhr",orgdb,rna=True)
        os.remove(aaseqs)
        os.remove(naseqs)
        checkpoint = "w2-4"
        log.info("JOB_CHECKPOINT::%s"%checkpoint)

    mlstdir = os.path.realpath(os.path.join(resultdir,"mlstgenes"))
    if not os.path.exists(mlstdir):
        os.makedirs(mlstdir)
    aligndir = os.path.realpath(os.path.join(resultdir,"mlst_aligned"))
    if not os.path.exists(aligndir):
        os.makedirs(aligndir)
    trimdir = os.path.realpath(os.path.join(resultdir,"mlst_trimmed"))
    if not os.path.exists(trimdir):
        os.makedirs(trimdir)

    # extract single copy results from genelist
    if checkpoint == "w2-4":
        genelist = [os.path.splitext(os.path.split(x)[1])[0] for x in glob.glob(os.path.join(refdir,reference,"*.fna"))]
        log.info("JOB_STATUS:: Writing MLST genes...")
        singles = [str(x) for x in getgenematrix.getmat(orgdb,rna=True,bh=True)[2][0] if x in genelist]
        getgenes.writeallgenes(orgdb,singles,[],outdir=mlstdir,outgroups=None,pct=0.5,writeaa=False,rename=True)

        log.info("JOB_STATUS:: Adding MLST genes to alignments...")
        addallalign([os.path.join(mlstdir,x+".fna") for x in singles], os.path.join(refdir,reference), aligndir, cpu=cpu,fast=fast)

        log.info("JOB_STATUS:: Trimming all alignments...")
        processmlst(aligndir,trimdir,cpu=cpu,trim=True)
        checkpoint = "w2-5"
        log.info("JOB_CHECKPOINT::%s"%checkpoint)

    treedir = os.path.realpath(os.path.join(resultdir,"trees"))
    if not os.path.exists(treedir):
        os.makedirs(treedir)

    #Add to prebuilt trees
    if checkpoint == "w2-5":
        log.info("JOB_PROGRESS::55/100")
        log.info("JOB_STATUS:: Adding to reference gene trees...")
        inlist = [os.path.join(trimdir,x) for x in os.listdir(trimdir)]
        addalltrees(inlist,os.path.join(refdir,reference),treedir,cpu=cpu)
        checkpoint = "w2-6"
        log.info("JOB_CHECKPOINT::%s" % checkpoint)

    if checkpoint == "w2-6":
        log.info("JOB_PROGRESS::85/100")
        log.info("JOB_STATUS:: Running coalescent tree phylogeny")
        #Combine all trees using ASTRAL
        flist = glob.glob(os.path.join(treedir,"RAxML_labelledTree.*.tree"))
        #reformat newick trees for astral
        for pbt in flist:
            with open(pbt,"r") as ifil, open(pbt+".newick","w") as ofil:
                x = ifil.next()
                ofil.write(re.sub("\[I\d+?\]|\"|'|QUERY___","",x))
        flist = glob.glob(os.path.join(treedir,"*.newick"))

        alltrees = catTrees(flist,os.path.join(treedir,"alltrees.tree"))
        coltree = os.path.join(treedir,"summaryTree.tree")
        runAstral(resultdir,alltrees,coltree)

        checkpoint = "w2-F"
        log.info("JOB_CHECKPOINT::%s" % checkpoint)

    finishedtree=""
    if checkpoint == "w2-F":
        finaltree = os.path.join(resultdir,"final.tree")
        if os.path.exists(os.path.join(treedir,"summaryTree.tree")):
            finishedtree = os.path.join(treedir,"summaryTree.tree")
        #Copy final tree to root dir
        if finishedtree:
            fmat = 2
            log.debug("Saving final tree... %s"%fmat)
            if not ete3helper.rerootTree(finishedtree,finaltree,fmat=fmat):
                shutil.copy(finishedtree,finaltree)
        else:
            log.error("Could not find final tree. Tree building failed")

def startjob(indir,resultdir,skip="",checkpoint=False,workflow=1,refdb="",cpu=1,concat=False,
             model="GTR",bs=0,kf=False,maxmlst=100,maxorg=50,filtMLST=True,refdir="",fast=False,minmlst=10):
    #Setup working directory
    if not os.path.exists(os.path.join(os.path.realpath(resultdir),"queryseqs")):
        os.makedirs(os.path.join(os.path.realpath(resultdir),"queryseqs")) #query sequence folder

    #Read checkpoint from log if exists or use checkpoint parameter
    if not checkpoint and os.path.exists(os.path.join(resultdir,"automlst.log")):
        with open(os.path.join(resultdir,"automlst.log"),"r") as fil:
            for line in fil:
                x = line.strip().split("::")
                if "JOB_CHECKPOINT" in x[0]:
                    checkpoint = x[1]
                elif "WORKFLOW" in x[0]:
                    workflow = int(x[1])
    #Start log
    global log
    log = setlog.init(os.path.join(resultdir,"automlst.log"),toconsole=True)

    #ensure bootsrtap value is valid and no more than 1000
    try:
        bs = int(bs)
        if bs > 1000:
            bs = 1000
        elif bs < 0:
            bs = 0
    except ValueError:
        log.warning("Invalid bootsrap value found, setting to 0")
        bs = 0

    log.debug("START / RESUME JOB last checkpoint: %s"%checkpoint)
    log.info('JOB_PARAMS::{"resultdir":"%s","skip":"%s","workflow":%s,"concat":"%s","model":"%s"}'%(resultdir,skip,workflow,concat,model))
    if workflow == 1:
        log.info("WORKFLOW::1")
        try:
            return startwf1(indir,resultdir,checkpoint=checkpoint,skip=skip,refdb=refdb,cpu=cpu,concat=concat,
                            model=model,bs=bs,kf=kf,maxmlst=maxmlst,maxorg=maxorg,filtMLST=filtMLST,fast=fast,minmlst=minmlst)
        except Exception as e:
            log.error("Unexpected failure: %s"%e)
            raise
    elif workflow == 2:
        log.info("WORKFLOW::2")
        try:
            return startwf2(indir,resultdir,refdir=refdir,checkpoint=checkpoint,model=model,bs=bs,kf=kf,maxmlst=maxmlst,cpu=cpu,fast=fast)
        except Exception as e:
            log.error("Unexpected failure: %s"%e)
            raise
    else:
        log.error("Improper workflow specified: %s"%workflow)
        return False

# Commandline Execution
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="""Start from input directory of genomes, get all reference distances and build phylogeny""")
    parser.add_argument("indir", help="Directory of input genomes or list of files")
    parser.add_argument("resultdir", help="Directory to store all results")
    parser.add_argument("-skip","--skip", help="Flag to skip manual intervention of organism selection. Ex: 'skip2.skip3' skips organism and mlst selection (default for cmdline execution)",default="skip2.skip3")
    parser.add_argument("-ref","--refdb", help="Reference database of orgs",default="")
    parser.add_argument("-rd","--refdirectory", help="Reference directory of pre-built trees and alignments",default="")
    parser.add_argument("-bs","--bootstrap", help="Bootstrap replicates (default: None)",default="")
    parser.add_argument("-cat","--concat", help="Use concatenated supermatrix to build tree",action="store_true",default=False)
    parser.add_argument("-kf","--keepfiles", help="Do not delete intermediate files after completion",action="store_true",default=False)
    parser.add_argument("-fm","--filtmlst", help="Filter mlst genes by starting tree congruance",action="store_true",default=False)
    parser.add_argument("-fa","--fastalign", help="Align using FFt-2 mode (faster)",action="store_true",default=False)
    parser.add_argument("-m","--model", help="Use specific model for iqtree parameter (default: GTR)",default="GTR")
    parser.add_argument("-ch","--checkpoint", help="Explicitly start at checkpoint",default=False)
    parser.add_argument("-c","--cpu", help="Number of cpu cores to use",type=int, default=1)
    parser.add_argument("-wf","--workflow", help="Workflow to use 1 or 2",choices=(1,2),type=int, default=1)
    parser.add_argument("-mm","--maxmlst", help="Maximum number of MLST genes to use (default:100)",type=int, default=100)
    parser.add_argument("-mn","--minmlst", help="Minimum number of MLST genes to use (default:10)",type=int, default=10)
    parser.add_argument("-mo","--maxorg", help="Maximum number of organisms to use (default:50)",type=int, default=50)
    args = parser.parse_args()
    startjob(args.indir,args.resultdir,args.skip,refdb=args.refdb,cpu=args.cpu,
             concat=args.concat,model=args.model,checkpoint=args.checkpoint,
             bs=args.bootstrap,kf=args.keepfiles,maxmlst=args.maxmlst,
             filtMLST=args.filtmlst,refdir=args.refdirectory,workflow=args.workflow,
             fast=args.fastalign,minmlst=args.minmlst)


# -*- coding: utf-8 -*-
#
# Copyright (c) 2015 Red Hat
# Licensed under The MIT License (MIT)
# http://opensource.org/licenses/MIT
#
from requests_kerberos.kerberos_ import log, _negotiate_value, HTTPKerberosAuth

def handle_401(self, response, **kwargs):
    """Handles 401's, attempts to use gssapi/kerberos authentication"""

    log.debug("handle_401(): Handling: 401")
    if _negotiate_value(response) is not None:
        _r = self.authenticate_user(response, **kwargs)
        if _r is response:
            log.debug("handle_401(): fail to authenticate, something wrong "
                      "with the client.")
            _r.raise_for_status()
        log.debug("handle_401(): returning {0}".format(_r))
        return _r
    else:
        log.debug("handle_401(): Kerberos is not supported")
        log.debug("handle_401(): returning {0}".format(response))
        return response

def authenticate_user(self, response, **kwargs):
    """Handles user authentication with gssapi/kerberos"""

    auth_header = self.generate_request_header(response)
    if auth_header is None:
        # GSS Failure, return existing response
        return response

    log.debug("authenticate_user(): Authorization header: {0}".format(
        auth_header))
    response.request.headers['Authorization'] = auth_header

    # Consume the content so we can reuse the connection for the next
    # request.
    response.content
    response.raw.release_conn()

    _r = response.connection.send(response.request, **kwargs)
    _r.history.append(response)

    log.debug("authenticate_user(): returning {0}".format(_r))
    return _r

def monkey_patch_kerberos():
    HTTPKerberosAuth.handle_401 = handle_401
    HTTPKerberosAuth.authenticate_user = authenticate_user