# coding=utf-8
"""Text Classifier based on bert."""
from __future__ import absolute_import, division, print_function

import collections
import csv
import gc
import io
import logging
import math
import multiprocessing
import os
import pickle
import time
from multiprocessing import Pool, cpu_count
from multiprocessing.dummy import Pool as ThreadPool
from urllib.parse import urlparse

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import PIL
import seaborn as sns
import sklearn
import tensorflow as tf
from sklearn.metrics import (
    classification_report, confusion_matrix, f1_score, precision_score,
    recall_score, roc_auc_score)
from tqdm import tqdm

import modeling
import optimization
import tokenization
from saver import ModelSaver
from tokenization import _is_punctuation

n_jobs = cpu_count()

class BertClassifier:
    def __init__(self, data_processor, num_labels, bert_config_file, max_seq_length, vocab_file, logdir, init_checkpoint, keep_checkpoint_max, use_GPU=False, label_smoothing=0.0, cycle=1):
        config = tf.ConfigProto(allow_soft_placement=True)
        # config.gpu_options.allow_growth = True
        self.sess = tf.Session(config=config)
        self.output_dropout_keep_prob = np.array([0.9])
        self.hidden_dropout_prob = np.array([0.1])
        self.attention_probs_dropout_prob = np.array([0.1])
        self.init_checkpoint = init_checkpoint
        bert_config = modeling.BertConfig.from_json_file(bert_config_file)
        # self.train_op, self.loss, self.logits, self.probabilities, self.feed_dict, self.attention_probs = create_model(
        self.train_op, self.loss, self.logits, self.probabilities, self.feed_dict = create_model(
            bert_config, num_labels, max_seq_length, self.sess, init_checkpoint=self.init_checkpoint, use_GPU=use_GPU, label_smoothing=label_smoothing, cycle=cycle)
        self.max_seq_length = max_seq_length
        self.tokenizer = tokenization.FullTokenizer(
            vocab_file=vocab_file, do_lower_case=True)
        if not os.path.exists(logdir):
            os.makedirs(logdir)
        self.summary_writer = tf.summary.FileWriter(logdir, self.sess.graph)
        self.prob_hist = None
        self.logits_hist = None
        self.eval_iterator = None
        self.num_eval_steps = None
        self.num_labels = num_labels
        self.model_saver = ModelSaver(keep_checkpoint_max=keep_checkpoint_max)
        self.data_processor = data_processor

    def train(self, data_dir, epochs, train_batch_size, eval_batch_size, learning_rate, warmup_proportion, save_checkpoints_steps, save_checkpoints_dir):
        '''Model train and eval.'''
        next_batch, num_train_steps, num_warmup_steps = create_data_iterator(self.data_processor,
                                                                             "train", data_dir, self.tokenizer, train_batch_size, self.max_seq_length, epochs, warmup_proportion)

        summary_ops = tf.summary.merge_all()
        # 断点训练
        ckpt = tf.train.get_checkpoint_state(save_checkpoints_dir)
        if ckpt and ckpt.model_checkpoint_path:
            tf.logging.info("Reload training state" +
                            (ckpt.model_checkpoint_path))
            self.model_saver.saver.restore(
                self.sess, ckpt.model_checkpoint_path)

        for step in range(num_train_steps):
            try:
                start = time.time()
                data = self.sess.run(next_batch)
                _, loss, global_step, merged_summary = self.sess.run((self.train_op, self.loss, tf.train.get_global_step(), summary_ops), feed_dict={self.feed_dict['input_ids']: data['input_ids'],
                                                                                self.feed_dict['input_mask']: data['input_mask'],
                                                                                self.feed_dict['segment_ids']: data['segment_ids'],
                                                                                self.feed_dict['label_ids']: data['label_ids'],
                                                                                self.feed_dict['sample_weight']: data['sample_weight'],
                                                                                self.feed_dict['output_dropout_keep_prob']: self.output_dropout_keep_prob,
                                                                                self.feed_dict['hidden_dropout_prob']: self.hidden_dropout_prob,
                                                                                self.feed_dict['attention_probs_dropout_prob']: self.attention_probs_dropout_prob,
                                                                                self.feed_dict['learning_rate']: learning_rate,
                                                                                self.feed_dict['num_train_steps']: num_train_steps,
                                                                                self.feed_dict['num_warmup_steps']: num_warmup_steps,
                                                                                self.feed_dict['batch_size']: train_batch_size})
                summary = tf.Summary(value=[tf.Summary.Value(
                    tag="Loss/Train", simple_value=loss)])
                self.summary_writer.add_summary(
                    summary, global_step=global_step)
                self.summary_writer.add_summary(
                    merged_summary, global_step=global_step)
                end = time.time()
                tf.logging.info("[%.2f%%] step: %d\tloss: %f\tcost time: %.3f" % (
                    (global_step/num_train_steps)*100, global_step, loss, (end-start)))
                if global_step % save_checkpoints_steps == 0 and global_step != 0:
                    fscore, auc, precision, recall = self.eval(
                        data_dir, eval_batch_size, is_training=True, global_step=global_step)
                    # 优先队列，保存前10个最好的checkpoints，之后可以做参数平均融合
                    self.model_saver.check_and_save_model(
                        save_checkpoints_dir, fscore, self.sess)
                if global_step > num_train_steps:
                    break
            except tf.errors.OutOfRangeError:
                break
        tf.logging.info("Train Finished.")
        # self.summary_writer.close()

    def eval(self, data_dir, eval_batch_size, is_training=False, global_step=None):
        if self.prob_hist == None:
            self.prob_hist = tf.summary.histogram(
                'prob_hist', self.probabilities)
        if self.logits_hist == None:
            self.logits_hist = tf.summary.histogram('logits_hist', self.logits)

        if (not is_training) or self.eval_iterator == None:
            self.eval_iterator, self.num_eval_steps, self.label_name = create_data_iterator(self.data_processor,
                                                                                "eval", data_dir, self.tokenizer, eval_batch_size, self.max_seq_length)
        self.sess.run(self.eval_iterator.initializer)
        loss_acc = []
        label_acc = None
        prob_acc = None
        start = time.time()
        for _ in tqdm(range(self.num_eval_steps), desc="Evaluation:"):
            try:
                data = self.sess.run(self.eval_iterator.get_next())
                loss, prob, prob_hist, logits_hist = self.sess.run((self.loss, self.probabilities, self.prob_hist, self.logits_hist), feed_dict={self.feed_dict['input_ids']: data['input_ids'],
                                                                                self.feed_dict['input_mask']: data['input_mask'],
                                                                                self.feed_dict['segment_ids']: data['segment_ids'],
                                                                                self.feed_dict['label_ids']: data['label_ids'],
                                                                                self.feed_dict['sample_weight']: data['sample_weight'],
                                                                                self.feed_dict['output_dropout_keep_prob']: np.array([1.0]),
                                                                                self.feed_dict['hidden_dropout_prob']: np.array([0.0]),
                                                                                self.feed_dict['attention_probs_dropout_prob']: np.array([0.0]),
                                                                                self.feed_dict['batch_size']: eval_batch_size})

                if isinstance(label_acc, type(None)):
                    assert loss_acc == [] and prob_acc == None
                    loss_acc.append(loss)
                    label_acc = data['label_ids']
                    prob_acc = prob
                else:
                    loss_acc.append(loss)
                    label_acc = np.concatenate(
                        (label_acc, data['label_ids']), axis=0)
                    prob_acc = np.concatenate((prob_acc, prob), axis=0)
            except tf.errors.OutOfRangeError:
                break

        assert len(prob_acc) == len(label_acc)

        # Classification report
        report = classification_report(label_acc, np.argmax(prob_acc, axis=-1), labels=[
                                       i for i in range(len(self.label_name))], target_names=self.label_name)
        tf.logging.info("***** Classification Report *****")
        tf.logging.info(report)
        # f1 score
        fscore = f1_score(label_acc, np.argmax(
            prob_acc, axis=-1), average='macro')
        # precision
        precision = precision_score(label_acc, np.argmax(
            prob_acc, axis=-1), average='macro')
        # recall
        recall = recall_score(label_acc, np.argmax(
            prob_acc, axis=-1), average='macro')
        # AUC
        auc = roc_auc_score(np.eye(self.num_labels)[
                            label_acc], prob_acc, average='macro')

        roc_curve, confusion_matrix = draw_image(prob_acc, np.eye(
            self.num_labels)[label_acc], label_acc, np.argmax(prob_acc, axis=-1), self.label_name)
        if is_training:
            summary = tf.Summary(value=[tf.Summary.Value(tag="Loss/Eval", simple_value=np.mean(loss_acc)),
                                        tf.Summary.Value(
                                            tag="Eval/auc", simple_value=auc),
                                        tf.Summary.Value(
                                            tag="Eval/f1_score", simple_value=fscore),
                                        tf.Summary.Value(
                                            tag="Eval/precision", simple_value=precision),
                                        tf.Summary.Value(
                                            tag="Eval/recall", simple_value=recall),
                                        tf.Summary.Value(tag='Eval_ROC', image=tf.Summary.Image(
                                            encoded_image_string=roc_curve)),
                                        tf.Summary.Value(tag='Eval_Confusion_Matrix', image=tf.Summary.Image(encoded_image_string=confusion_matrix))])
            self.summary_writer.add_summary(
                prob_hist, global_step=global_step)
            self.summary_writer.add_summary(
                logits_hist, global_step=global_step)
            self.summary_writer.add_summary(
                summary, global_step=global_step)
        end = time.time()
        tf.logging.info("Evaluation Finished.\tcost time: %.3f\tF1 Score:%.3f\tAuc:%.3f\tprecision:%.3f\trecall:%.3f\t" % (
            (end-start), fscore, auc, precision, recall))

        return fscore, auc, precision, recall

    def predict(self, predict_batch_size=1, output_dir='./predict', file_path=None, input_example=None):
        # print((file_path != None or input_example != None))
        assert (file_path != None or input_example !=
                None), "file_path and input_example must have one not None."
        if file_path != None:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            out_file = open(os.path.join(output_dir, "out.txt"),
                            'w', encoding="utf-8")
            next_batch, num_predict_steps, label_name = create_data_iterator(self.data_processor,
                                                                             "predict", file_path, self.tokenizer, predict_batch_size, self.max_seq_length)
            out_file.write("prob,label\n")
            for _ in tqdm(range(num_predict_steps), desc="Predicting:"):
                try:
                    data = self.sess.run(next_batch)
                    prob = self.sess.run((self.probabilities), feed_dict={self.feed_dict['input_ids']: data['input_ids'],
                                                                          self.feed_dict['input_mask']: data['input_mask'],
                                                                          self.feed_dict['segment_ids']: data['segment_ids'],
                                                                          self.feed_dict['output_dropout_keep_prob']: np.array([1.0]),
                                                                          self.feed_dict['hidden_dropout_prob']: np.array([0.0]),
                                                                          self.feed_dict['attention_probs_dropout_prob']: np.array([0.0]),
                                                                          self.feed_dict['batch_size']: predict_batch_size})
                    for p in prob:
                        out_file.write("%s,%s\n" %
                                       (p, label_name[np.argmax(p)]))
                except tf.errors.OutOfRangeError:
                    break
            out_file.close()
        else:
            s = time.time()
            input_feature = convert_single_example(
                0, input_example, None, self.max_seq_length, self.tokenizer, is_predict=True)
            e = time.time()
            print("process:", e-s)
            s = time.time()
            prob = self.sess.run((self.probabilities), feed_dict={self.feed_dict['input_ids']: np.expand_dims(np.array(input_feature.input_ids), axis=0),
                                                                  self.feed_dict['input_mask']: np.expand_dims(np.array(input_feature.input_mask), axis=0),
                                                                  self.feed_dict['segment_ids']: np.expand_dims(np.array(input_feature.segment_ids), axis=0),
                                                                  self.feed_dict['output_dropout_keep_prob']: np.array([1.0]),
                                                                  self.feed_dict['hidden_dropout_prob']: np.array([0.0]),
                                                                  self.feed_dict['attention_probs_dropout_prob']: np.array([0.0]),
                                                                  self.feed_dict['batch_size']: predict_batch_size})
            e = time.time()
            print("inference:", e-s)
            return prob

class InputExample(object):
    """A single training/test example for simple sequence classification."""

    def __init__(self, guid, text_a, text_b=None, text_c=None, label=None, weight=1):
        """Constructs a InputExample.

        Args:
          guid: Unique id for the example.
          text_a: string. The untokenized text of the first sequence. For single
            sequence tasks, only this sequence must be specified.
          text_b: (Optional) string. The untokenized text of the second sequence.
            Only must be specified for sequence pair tasks.
          text_c: (Optional) string. The untokenized text of the third sequence.
            Only must be specified for sequence pair tasks.
          label: (Optional) string. The label of the example. This should be
            specified for train and dev examples, but not for test examples.
          weight: (Optional) float. The weight of the example. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text_a = text_a
        self.text_b = text_b
        self.text_c = text_c
        self.label = label
        self.weight = weight

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self,
                 input_ids,
                 input_mask,
                 segment_ids,
                 label_id,
                 weight):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_id = label_id
        self.sample_weight = weight

def create_data_iterator(processor, mode, data_dir, tokenizer, batch_size, max_seq_length=192, epochs=1, warmup_proportion=None):
    label_list = processor.get_labels()
    if mode == "train":
        # 如果数据集存在，则不用导入数据，直接训练
        if os.path.exists(os.path.join(data_dir, "train.tf_record_0")):
            train_examples_cnt = 0
            with open(os.path.join(data_dir, "train.tsv"), 'r', encoding='utf-8') as f:
                for _ in f:
                    train_examples_cnt += 1
            num_train_steps = math.ceil(
                train_examples_cnt / batch_size * epochs)
            num_warmup_steps = int(num_train_steps * warmup_proportion)
            # print log
            tf.logging.info("***** Running training *****")
            tf.logging.info("  Num examples = %d", train_examples_cnt)
            tf.logging.info("  Batch size = %d", batch_size)
            tf.logging.info("  Num steps = %d", num_train_steps)
        else:
            train_examples = processor.get_train_examples(data_dir)
            num_train_steps = math.ceil(
                len(train_examples) / batch_size * epochs)
            num_warmup_steps = int(num_train_steps * warmup_proportion)
            # print log
            tf.logging.info("***** Running training *****")
            tf.logging.info("  Num examples = %d", len(train_examples))
            tf.logging.info("  Batch size = %d", batch_size)
            tf.logging.info("  Num steps = %d", num_train_steps)
            # convert examples to features and write to tf-record files.
            file_based_convert_examples_to_features(
                train_examples, label_list, max_seq_length, tokenizer, os.path.join(data_dir, "train.tf_record"))

        # Load multiple tf-record dateset files.
        filenames = [os.path.join(data_dir, "train.tf_record_%d" % i)
                     for i in range(n_jobs)]
        data_set = tf.data.TFRecordDataset(filenames)
        # Create a description of the features.
        feature_description = {
            'input_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64, default_value=[0 for i in range(max_seq_length)]),
            'input_mask': tf.io.FixedLenFeature([max_seq_length], tf.int64, default_value=[0 for i in range(max_seq_length)]),
            'segment_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64, default_value=[0 for i in range(max_seq_length)]),
            'label_ids': tf.io.FixedLenFeature([], tf.int64, default_value=0),
            'sample_weight': tf.io.FixedLenFeature([], tf.float32, default_value=0.0),
        }

        def _parse_function(example_proto):
            # Parse the input `tf.Example` proto using the dictionary above.
            return tf.io.parse_single_example(example_proto, feature_description)

        data_set = data_set.map(_parse_function)
        data_set = data_set.repeat(int(epochs))
        data_set = data_set.shuffle(buffer_size=100)

        data_set = data_set.batch(batch_size=batch_size)
        iterator = data_set.make_one_shot_iterator()
        next_batch = iterator.get_next()

        return next_batch, num_train_steps, num_warmup_steps

    elif mode == "eval":
        eval_examples = processor.get_eval_examples(data_dir)
        num_eval_steps = math.ceil(len(eval_examples) / batch_size)

        # print log
        tf.logging.info("***** Running evaluation *****")
        tf.logging.info("  Num examples = %d ", len(eval_examples))
        tf.logging.info("  Batch size = %d", batch_size)
        tf.logging.info("  Num steps = %d", num_eval_steps)

        features = convert_examples_to_features(
            eval_examples, label_list, max_seq_length, tokenizer)

        all_input_ids = []
        all_input_mask = []
        all_segment_ids = []
        all_label_ids = []
        all_sample_weight = []

        for feature in features:
            all_input_ids.append(feature.input_ids)
            all_input_mask.append(feature.input_mask)
            all_segment_ids.append(feature.segment_ids)
            all_label_ids.append(feature.label_id)
            all_sample_weight.append(feature.sample_weight)

        num_examples = len(features)

        data_set = tf.data.Dataset.from_tensor_slices({
            "input_ids":
                tf.constant(
                    all_input_ids, shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
            "input_mask":
                tf.constant(
                    all_input_mask,
                    shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
            "segment_ids":
                tf.constant(
                    all_segment_ids,
                    shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
            "label_ids":
                tf.constant(all_label_ids, shape=[
                            num_examples], dtype=tf.int32),
            "sample_weight":
                tf.constant(all_sample_weight, shape=[
                            num_examples], dtype=tf.float32),
        })

        data_set = data_set.batch(batch_size=batch_size)
        iterator = data_set.make_initializable_iterator()
        # next_batch = iterator.get_next()

        return iterator, num_eval_steps, label_list

    elif mode == "predict":
        predict_examples = processor.get_test_examples(data_dir)
        num_predict_steps = math.ceil(len(predict_examples) / batch_size)

        # print log
        tf.logging.info("***** Running predict *****")
        tf.logging.info("  Num examples = %d ", len(predict_examples))
        tf.logging.info("  Batch size = %d", batch_size)
        tf.logging.info("  Num steps = %d", num_predict_steps)

        features = convert_examples_to_features(
            predict_examples, label_list, max_seq_length, tokenizer)

        all_input_ids = []
        all_input_mask = []
        all_segment_ids = []

        for feature in features:
            all_input_ids.append(feature.input_ids)
            all_input_mask.append(feature.input_mask)
            all_segment_ids.append(feature.segment_ids)

        num_examples = len(features)

        data_set = tf.data.Dataset.from_tensor_slices({
            "input_ids":
                tf.constant(
                    all_input_ids, shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
            "input_mask":
                tf.constant(
                    all_input_mask,
                    shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
            "segment_ids":
                tf.constant(
                    all_segment_ids,
                    shape=[num_examples, max_seq_length],
                    dtype=tf.int32),
        })

        data_set = data_set.batch(batch_size=batch_size)
        iterator = data_set.make_one_shot_iterator()
        next_batch = iterator.get_next()

        return next_batch, num_predict_steps, label_list

    else:
        raise ValueError("Mode must be 'train', 'eval' or 'predict'.")

def plot_roc_curve(prob, label_onehot):
    fpr, tpr, thresholds = sklearn.metrics.roc_curve(
        label_onehot.ravel(), prob.ravel())
    auc = sklearn.metrics.auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, c='r', lw=2, alpha=0.7,
             label=u'AUC=%.3f' % auc)
    plt.plot((0, 1), (0, 1), c='#808080', lw=1, ls='--', alpha=0.7)
    plt.xlim((-0.01, 1.02))
    plt.ylim((-0.01, 1.02))
    plt.xticks(np.arange(0, 1.1, 0.1))
    plt.yticks(np.arange(0, 1.1, 0.1))
    plt.xlabel('False Positive Rate', fontsize=13)
    plt.ylabel('True Positive Rate', fontsize=13)
    plt.grid(b=True, ls=':')
    plt.legend(loc='lower right', fancybox=True,
               framealpha=0.8, fontsize=12)
    plt.title(u'Eval ROC And AUC', fontsize=17)
    buffer_ = io.BytesIO()
    plt.savefig(buffer_, format='png')
    buffer_.seek(0)
    png_string = buffer_.getvalue()
    buffer_.close()
    plt.close()
    return png_string

def plot_confusion_matrix(label_id, predicted_classes, label_names):
    # Compute confusion matrix
    cm = sklearn.metrics.confusion_matrix(
        label_id, predicted_classes)
    plt.figure()
    sns.heatmap(cm, annot=True, yticklabels=label_names,
                xticklabels=label_names, linewidths=.5)
    plt.yticks(rotation=360)
    plt.title(u'Confusion Matrix Heat Map', fontsize=17)
    plt.tight_layout()
    buffer_ = io.BytesIO()
    plt.savefig(buffer_, format='png')
    buffer_.seek(0)
    png_string = buffer_.getvalue()
    buffer_.close()
    plt.close()

    return png_string

def draw_image(prob, label_onehot, label_id, predicted_classes, label_names):
    roc_curve = plot_roc_curve(prob, label_onehot)
    confusion_matrix = plot_confusion_matrix(
        label_id, predicted_classes, label_names)

    return roc_curve, confusion_matrix

def process_url(url):
    """Converts a string url to a list of token string."""
    # only get url path, remove host,params.
    url = urlparse(url).path
    # url = list(url)
    # for i in range(len(url)):
    #     if _is_punctuation(url[i]):
    #         url[i] = " "
    # url = ''.join(url)
    # url = ' '.join(url.split())
    return url

def convert_single_example(ex_index, example, label_list, max_seq_length,
                           tokenizer, is_predict=False):
    """Converts a single `InputExample` into a single `InputFeatures`."""
    # title
    tokens_a = tokenizer.tokenize(example.text_a)
    tokens_b = None
    tokens_c = None
    if example.text_b:
        # URL pre-process
        url = process_url(example.text_b)
        tokens_b = tokenizer.tokenize(url)

    if example.text_c:
        # body
        tokens_c = tokenizer.tokenize(example.text_c)

    if tokens_b and tokens_c:
        # Modifies `tokens_a`, `tokens_b` and `tokens_c` in place so that the total
        # length is less than the specified length.
        # Account for [CLS], [SEP], [SEP], [SEP]with "- 4"
        _truncate_seq_pair_3(tokens_a, tokens_b, tokens_c, max_seq_length - 4)
    elif tokens_b:
         # Modifies `tokens_a` and `tokens_b` in place so that the total
        # length is less than the specified length.
        # Account for [CLS], [SEP], [SEP] with "- 3"
        _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)
    else:
        # Account for [CLS] and [SEP] with "- 2"
        if len(tokens_a) > max_seq_length - 2:
            tokens_a = tokens_a[0:(max_seq_length - 2)]

    # The convention in BERT is:
    # (a) For sequence pairs:
    #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
    #  type_ids: 0     0  0    0    0     0       0 0     1  1  1  1   1 1
    # (b) For single sequences:
    #  tokens:   [CLS] the dog is hairy . [SEP]
    #  type_ids: 0     0   0   0  0     0 0
    #
    # Where "type_ids" are used to indicate whether this is the first
    # sequence or the second sequence. The embedding vectors for `type=0` and
    # `type=1` were learned during pre-training and are added to the wordpiece
    # embedding vector (and position vector). This is not *strictly* necessary
    # since the [SEP] token unambiguously separates the sequences, but it makes
    # it easier for the model to learn the concept of sequences.
    #
    # For classification tasks, the first vector (corresponding to [CLS]) is
    # used as the "sentence vector". Note that this only makes sense because
    # the entire model is fine-tuned.
    tokens = []
    segment_ids = []
    tokens.append("[CLS]")
    segment_ids.append(0)
    for token in tokens_a:
        tokens.append(token)
        segment_ids.append(0)
    tokens.append("[SEP]")
    segment_ids.append(0)

    if tokens_b:
        for token in tokens_b:
            tokens.append(token)
            segment_ids.append(1)
        tokens.append("[SEP]")
        segment_ids.append(1)

    if tokens_c:
        for token in tokens_c:
            tokens.append(token)
            segment_ids.append(2)
        tokens.append("[SEP]")
        segment_ids.append(2)

    input_ids = tokenizer.convert_tokens_to_ids(tokens)

    # The mask has 1 for real tokens and 0 for padding tokens. Only real
    # tokens are attended to.
    input_mask = [1] * len(input_ids)

    # Zero-pad up to the sequence length.
    while len(input_ids) < max_seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)

    assert len(input_ids) == max_seq_length
    assert len(input_mask) == max_seq_length
    assert len(segment_ids) == max_seq_length

    if is_predict:
        feature = InputFeatures(
            input_ids=input_ids,
            input_mask=input_mask,
            segment_ids=segment_ids,
            label_id=None,
            weight=None)
        return feature

    label_map = {}
    for (i, label) in enumerate(label_list):
        label_map[label] = i

    label_id = label_map[example.label]
    # if ex_index < 5:
    #     print("*** Example ***")
    #     print("guid: %s" % (example.guid))
    #     print("tokens: %s" % " ".join(
    #         [tokenization.printable_text(x) for x in tokens]))
    #     print("input_ids: %s" %
    #                     " ".join([str(x) for x in input_ids]))
    #     print("input_mask: %s" %
    #                     " ".join([str(x) for x in input_mask]))
    #     print("segment_ids: %s" %
    #                     " ".join([str(x) for x in segment_ids]))
    #     print("label: %s (id = %d)" % (example.label, label_id))

    feature = InputFeatures(
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids,
        label_id=label_id,
        weight=example.weight)
    return feature

def file_based_convert(task_id, output_file, examples, label_list, max_seq_length, tokenizer):
    try:
        writer = tf.python_io.TFRecordWriter(output_file)
        for (ex_index, example) in enumerate(examples):
            if ex_index % 10000 == 0:
                print("Task %d\t:Writing example %d of %d" %
                      (task_id, ex_index, len(examples)))
            feature = convert_single_example(ex_index, example, label_list,
                                             max_seq_length, tokenizer)
            features = collections.OrderedDict()
            features["input_ids"] = tf.train.Feature(
                int64_list=tf.train.Int64List(value=list(feature.input_ids)))
            features["input_mask"] = tf.train.Feature(
                int64_list=tf.train.Int64List(value=list(feature.input_mask)))
            features["segment_ids"] = tf.train.Feature(
                int64_list=tf.train.Int64List(value=list(feature.segment_ids)))
            features["label_ids"] = tf.train.Feature(
                int64_list=tf.train.Int64List(value=list([feature.label_id])))
            features["sample_weight"] = tf.train.Feature(
                float_list=tf.train.FloatList(value=list([feature.sample_weight])))
            tf_example = tf.train.Example(
                features=tf.train.Features(feature=features))
            writer.write(tf_example.SerializeToString())
    except Exception as e:
        print(e)
    writer.close()

def file_based_convert_examples_to_features(
        examples, label_list, max_seq_length, tokenizer, output_file):
    """Convert a set of `InputExample`s to a TFRecord file."""
    if os.path.exists(output_file+"_0"):
        return

    p = Pool(n_jobs)
    chunk_size = int(len(examples) / n_jobs)

    for i in range(n_jobs):
        if i < n_jobs - 1:
            p.apply_async(file_based_convert, args=(
                i,  "%s_%d" % (output_file, i), examples[i*chunk_size: (i+1) * chunk_size], label_list, max_seq_length, tokenizer,))
        else:
            p.apply_async(file_based_convert, args=(
                i, "%s_%d" % (output_file, i), examples[i*chunk_size: len(examples)], label_list, max_seq_length, tokenizer,))
    p.close()
    p.join()

def convert_examples_to_features(examples, label_list, max_seq_length,
                                 tokenizer):
    """Convert a set of `InputExample`s to a list of `InputFeatures`."""

    features = []
    for (ex_index, example) in enumerate(examples):
        if ex_index % 10000 == 0:
            tf.logging.info("Writing example %d of %d" %
                            (ex_index, len(examples)))

        feature = convert_single_example(ex_index, example, label_list,
                                         max_seq_length, tokenizer)

        features.append(feature)
    return features

def _truncate_seq_pair_3(tokens_a, tokens_b, tokens_c, max_length):
    """Truncates a sequence pair in place to the maximum length."""

    # This is a simple heuristic which will always truncate the longer sequence
    # one token at a time. This makes more sense than truncating an equal percent
    # of tokens from each, since if one sequence is very short then each token
    # that's truncated likely contains more information than a longer sequence.
    while True:
        total_length = len(tokens_a) + len(tokens_b) + len(tokens_c)
        if total_length <= max_length:
            break
        if len(tokens_b) > len(tokens_c):
            tokens_b.pop()
        else:
            tokens_c.pop()

def _truncate_seq_pair(tokens_a, tokens_b, max_length):
    """Truncates a sequence pair in place to the maximum length."""

    # This is a simple heuristic which will always truncate the longer sequence
    # one token at a time. This makes more sense than truncating an equal percent
    # of tokens from each, since if one sequence is very short then each token
    # that's truncated likely contains more information than a longer sequence.
    while True:
        total_length = len(tokens_a) + len(tokens_b)
        if total_length <= max_length:
            break
        if len(tokens_a) > len(tokens_b):
            tokens_a.pop()
        else:
            tokens_b.pop()

def get_available_gpus():
    """
    code from http://stackoverflow.com/questions/38559755/how-to-get-current-available-gpus-in-tensorflow
    """
    from tensorflow.python.client import device_lib as _device_lib
    local_device_protos = _device_lib.list_local_devices()
    return [x.name.split(":")[-1] for x in local_device_protos if x.device_type == 'GPU']

def average_gradients(tower_grads):
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        grads = []
        for g, _ in grad_and_vars:
            expend_g = tf.expand_dims(g, 0)
            grads.append(expend_g)
        grad = tf.concat(grads, 0)
        grad = tf.reduce_mean(grad, 0)
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        average_grads.append(grad_and_var)
    return average_grads

def assign_to_device(device, ps_device='/cpu:0'):
    def _assign(op):
        node_def = op if isinstance(op, tf.NodeDef) else op.node_def
        if node_def.op in ['Variable', 'VariableV2', 'AutoReloadVariable']:
            return "/" + ps_device
        else:
            return device

    return _assign

def create_model(bert_config, num_labels, max_seq_length, sess, init_checkpoint=None, use_GPU=False, label_smoothing=0.0, cycle=1):
    """Creates a classification model."""
    GPUs = get_available_gpus()
    defalut_device = '/cpu:0'
    if use_GPU and len(GPUs) != 0:
        defalut_device = '/gpu:{}'.format(GPUs[0])
    # Place all ops on CPU by default
    with tf.device(defalut_device):
        tower_grads = []
        loss_list = []
        logits_list = []
        probabilities_list = []
        train_op = None
        loss = None
        logits = None
        probabilities = None
        global_step = tf.train.get_or_create_global_step()
        # input placeholder
        _input_ids = tf.placeholder(tf.int64, shape=(None, max_seq_length))
        _input_mask = tf.placeholder(tf.int64, shape=(None, max_seq_length))
        _segment_ids = tf.placeholder(tf.int64, shape=(None, max_seq_length))
        _label_ids = tf.placeholder(tf.int64, shape=None)
        _sample_weight = tf.placeholder(tf.float32, shape=None)
        _output_dropout_keep_prob = tf.placeholder(tf.float32, shape=None)
        _hidden_dropout_prob = tf.placeholder(tf.float32, shape=None)
        _attention_probs_dropout_prob = tf.placeholder(tf.float32, shape=None)
        # optimizer placeholder
        _learning_rate = tf.placeholder(tf.float32, shape=None)
        _num_train_steps = tf.placeholder(tf.int32, shape=None)
        _num_warmup_steps = tf.placeholder(tf.int32, shape=None)
        _batch_size = tf.placeholder(tf.int32, shape=None)
        # feed dict
        feed_dict = {'input_ids': _input_ids,
                     'input_mask': _input_mask,
                     'segment_ids': _segment_ids,
                     'label_ids': _label_ids,
                     'sample_weight': _sample_weight,
                     'output_dropout_keep_prob': _output_dropout_keep_prob,
                     'hidden_dropout_prob': _hidden_dropout_prob,
                     'attention_probs_dropout_prob': _attention_probs_dropout_prob,
                     'learning_rate': _learning_rate,
                     'num_train_steps': _num_train_steps,
                     'num_warmup_steps': _num_warmup_steps,
                     'batch_size': _batch_size}

        optimizer = optimization.create_optimizer(
            _learning_rate, tf.cast((_num_train_steps / cycle), tf.int32), _num_warmup_steps)
        if use_GPU:
            batch_size = tf.to_int32(_batch_size / len(GPUs))
            for i in range(len(GPUs)):
                # with tf.device(assign_to_device('/gpu:{}'.format(GPUs[i]), ps_device='/gpu:0')):
                with tf.device('/gpu:{}'.format(GPUs[i])):
                    # split input data for every gpu device.
                    with tf.name_scope("input_slice"):
                        input_ids = _input_ids[i *
                                               batch_size:(i + 1) * batch_size]
                        input_mask = _input_mask[i *
                                                 batch_size:(i + 1) * batch_size]
                        segment_ids = _segment_ids[i *
                                                   batch_size:(i + 1) * batch_size]
                        label_ids = _label_ids[i *
                                               batch_size:(i + 1) * batch_size]
                        sample_weight = _sample_weight[i *
                                                       batch_size:(i + 1) * batch_size]

                    # build model
                    model = modeling.BertModel(
                        config=bert_config,
                        input_ids=input_ids,
                        input_mask=input_mask,
                        token_type_ids=segment_ids,
                        hidden_dropout_prob=_hidden_dropout_prob,
                        attention_probs_dropout_prob=_attention_probs_dropout_prob,
                        scope="bert")
                    # If you want to use the token-level output, use model.get_sequence_output() instead.
                    output_layer = model.get_pooled_output()
                    hidden_size = output_layer.shape[-1].value
                    with tf.variable_scope("output", reuse=tf.AUTO_REUSE):
                        output_weights = tf.get_variable(
                            "output_weights", [num_labels, hidden_size],
                            initializer=tf.truncated_normal_initializer(stddev=0.02))

                        output_bias = tf.get_variable(
                            "output_bias", [num_labels], initializer=tf.zeros_initializer())

                        with tf.variable_scope("loss"):
                            # I.e., 0.1 dropout
                            output_layer = tf.nn.dropout(
                                output_layer, keep_prob=_output_dropout_keep_prob)

                            logits_ = tf.matmul(
                                output_layer, output_weights, transpose_b=True)
                            logits_ = tf.nn.bias_add(logits_, output_bias)
                            probabilities_ = tf.nn.softmax(logits_, axis=-1)

                            one_hot_labels = tf.one_hot(
                                label_ids, depth=num_labels, dtype=tf.float32)

                            loss_ = tf.losses.softmax_cross_entropy(
                                one_hot_labels,
                                logits_,
                                weights=sample_weight,
                                label_smoothing=label_smoothing
                            )

                            grads_ = optimizer.compute_gradients(loss_)
                            tower_grads.append(grads_)
                            loss_list.append(loss_)
                            logits_list.append(logits_)
                            probabilities_list.append(probabilities_)

            loss = tf.reduce_mean(loss_list)
            if len(GPUs) == 1:
                logits = tf.squeeze(logits_list, [0])
                probabilities = tf.squeeze(probabilities_list, [0])
            else:
                logits = tf.keras.layers.concatenate(logits_list, axis=0)
                probabilities = tf.keras.layers.concatenate(
                    probabilities_list, axis=0)
            # Merge grads
            with tf.name_scope("merge_grads"):
                grads = average_gradients(tower_grads)
                capped_gvs = [(tf.clip_by_value(grad, -1., 1.), var)
                              for grad, var in grads]
            train_op = optimizer.apply_gradients(
                capped_gvs, global_step=global_step)
        else:
            # build model
            model = modeling.BertModel(
                config=bert_config,
                input_ids=_input_ids,
                input_mask=_input_mask,
                token_type_ids=_segment_ids,
                hidden_dropout_prob=_hidden_dropout_prob,
                attention_probs_dropout_prob=_attention_probs_dropout_prob,
                scope="bert")
            # If you want to use the token-level output, use model.get_sequence_output() instead.
            output_layer = model.get_pooled_output()
            hidden_size = output_layer.shape[-1].value
            with tf.variable_scope("output", reuse=tf.AUTO_REUSE):
                output_weights = tf.get_variable(
                    "output_weights", [num_labels, hidden_size],
                    initializer=tf.truncated_normal_initializer(stddev=0.02))

                output_bias = tf.get_variable(
                    "output_bias", [num_labels], initializer=tf.zeros_initializer())

                with tf.variable_scope("loss"):
                    # I.e., 0.1 dropout
                    output_layer = tf.nn.dropout(
                        output_layer, keep_prob=_output_dropout_keep_prob)

                    logits = tf.matmul(
                        output_layer, output_weights, transpose_b=True)
                    logits = tf.nn.bias_add(logits, output_bias)
                    probabilities = tf.nn.softmax(logits, axis=-1)

                    one_hot_labels = tf.one_hot(
                        _label_ids, depth=num_labels, dtype=tf.float32)

                    loss = tf.losses.softmax_cross_entropy(
                        one_hot_labels,
                        logits,
                        weights=_sample_weight,
                        label_smoothing=label_smoothing
                    )
            with tf.name_scope("merge_grads"):
                grads = optimizer.compute_gradients(loss)
                capped_gvs = [(tf.clip_by_value(grad, -1., 1.), var)
                              for grad, var in grads]
            train_op = optimizer.apply_gradients(
                capped_gvs, global_step=global_step)

        # initial model's variables.
        tf.logging.info("Load model checkpoint : %s" % init_checkpoint)
        tvars = tf.trainable_variables()
        if init_checkpoint:
            (assignment_map, initialized_variable_names
             ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)

            tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

        init_op = tf.global_variables_initializer()
        sess.run(init_op)

        # # print variables
        # tf.logging.info("**** Trainable Variables ****")
        # for var in tvars:
        #     init_string = ""
        #     if var.name in initialized_variable_names:
        #         init_string = ", *INIT_FROM_CKPT*"
        #     tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
        #                     init_string)
        # attention_probs = model.get_all_layer_attention_probs()
        # return (train_op, loss, logits, probabilities, feed_dict, attention_probs)
        return (train_op, loss, logits, probabilities, feed_dict)


import argparse
import logging
import os
import socket
import sys

import numpy as np
import psutil
import setproctitle
import torch
import wandb
from mpi4py import MPI

# add the FedML root directory to the python path

sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../../")))
from fedml_api.data_preprocessing.cifar10.data_loader import load_partition_data_cifar10
from fedml_api.data_preprocessing.cifar100.data_loader import load_partition_data_cifar100
from fedml_api.data_preprocessing.cinic10.data_loader import load_partition_data_cinic10
from fedml_api.distributed.fedgkt.FedGKTAPI import FedML_init, FedML_FedGKT_distributed

from fedml_api.model.cv.resnet56_gkt.resnet_client import resnet8_56
from fedml_api.model.cv.resnet56_gkt.resnet_pretrained import resnet56_pretrained
from fedml_api.model.cv.resnet56_gkt.resnet_server import resnet56_server

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

def add_args(parser):
    """
    parser : argparse.ArgumentParser
    return a parser added with args required by fit
    """
    # Training settings
    parser.add_argument('--model_client', type=str, default='resnet5', metavar='N',
                        help='neural network used in training')

    parser.add_argument('--model_server', type=str, default='resnet32', metavar='N',
                        help='neural network used in training')

    parser.add_argument('--dataset', type=str, default='cifar10', metavar='N',
                        help='dataset used for training')

    parser.add_argument('--data_dir', type=str, default='./data', help='data directory')

    parser.add_argument('--partition_method', type=str, default='homo', metavar='N',
                        help='how to partition the dataset on local workers')

    parser.add_argument('--partition_alpha', type=float, default=0.5, metavar='PA',
                        help='partition alpha (default: 0.5)')

    parser.add_argument('--batch_size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')

    parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
                        help='learning rate (default: 0.001)')

    parser.add_argument('--wd', help='weight decay parameter;', type=float, default=5e-4)

    parser.add_argument('--epochs_client', type=int, default=3, metavar='EP',
                        help='how many epochs will be trained locally')

    parser.add_argument('--local_points', type=int, default=5000, metavar='LP',
                        help='the approximate fixed number of data points we will have on each local worker')

    parser.add_argument('--client_number', type=int, default=4, metavar='NN',
                        help='number of workers in a distributed cluster')

    parser.add_argument('--comm_round', type=int, default=300,
                        help='how many round of communications we shoud use')

    parser.add_argument('--gpu', type=int, default=0,
                        help='gpu')

    parser.add_argument('--loss_scale', type=float, default=1024,
                        help='Loss scaling, positive power of 2 values can improve fp16 convergence.')
    parser.add_argument('--no_bn_wd', action='store_true', help='Remove batch norm from weight decay')

    # knowledge distillation
    parser.add_argument('--temperature', default=3.0, type=float, help='Input the temperature: default(3.0)')
    parser.add_argument('--epochs_server', type=int, default=3, metavar='EP',
                        help='how many epochs will be trained on the server side')
    parser.add_argument('--alpha', default=1.0, type=float, help='Input the relative weight: default(1.0)')
    parser.add_argument('--optimizer', default="SGD", type=str, help='optimizer: SGD, Adam, etc.')
    parser.add_argument('--whether_training_on_client', default=1, type=int)
    parser.add_argument('--whether_distill_on_the_server', default=0, type=int)
    parser.add_argument('--client_model', default="resnet4", type=str)
    parser.add_argument('--weight_init_model', default="resnet32", type=str)
    parser.add_argument('--running_name', default="default", type=str)
    parser.add_argument('--sweep', default=0, type=int)
    parser.add_argument('--multi_gpu_server', action='store_true')
    parser.add_argument('--test', action='store_true',
                        help='test mode, only run 1-2 epochs to test the bug of the program')

    parser.add_argument('--gpu_num_per_server', type=int, default=8,
                        help='gpu_num_per_server')

    args = parser.parse_args()
    return args

def load_data(args, dataset_name):
    if dataset_name == "cifar10":
        data_loader = load_partition_data_cifar10
    elif dataset_name == "cifar100":
        data_loader = load_partition_data_cifar100
    elif dataset_name == "cinic10":
        data_loader = load_partition_data_cinic10
    else:
        data_loader = load_partition_data_cifar10

    train_data_num, test_data_num, train_data_global, test_data_global, \
    train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
    class_num = data_loader(args.dataset, args.data_dir, args.partition_method,
                            args.partition_alpha, args.client_number, args.batch_size)

    dataset = [train_data_num, test_data_num, train_data_global, test_data_global,
               train_data_local_num_dict, train_data_local_dict, test_data_local_dict, class_num]
    return dataset

def create_client_model(args, n_classes):
    """
        Note that we only initialize the client feature extractor to mitigate the difficulty of alternating optimization
    """
    if args.dataset == "cifar10" or args.dataset == "CIFAR10":
        resumePath = "./../../../fedml_api/model/cv/pretrained/CIFAR10/resnet56/best.pth"
    elif args.dataset == "cifar100" or args.dataset == "CIFAR100":
        resumePath = "./../../../fedml_api/model/cv/pretrained/CIFAR100/resnet56/best.pth"
    elif args.dataset == "cinic10" or args.dataset == "CINIC10":
        resumePath = "./../../../fedml_api/model/cv/pretrained/CINIC10/resnet56/best.pth"
    else:
        resumePath = "./../../../fedml_api/model/cv/pretrained/CIFAR10/resnet56/best.pth"
    pretrained_model = resnet56_pretrained(n_classes, pretrained=True, path=resumePath)
    logging.info("########pretrained model#################")
    # logging.info(pretrained_model)

    # copy pretrained parameters to client models
    params_featrue_extractor = dict()
    for name, param in pretrained_model.named_parameters():
        if name.startswith("conv1") or name.startswith("bn1") or name.startswith("layer1"):
            # logging.info(name)
            params_featrue_extractor[name] = param

    client_model = resnet8_56(n_classes)

    # logging.info("pretrained:")
    for name, param in client_model.named_parameters():
        if name.startswith("conv1"):
            param.data = params_featrue_extractor[name]
            if args.whether_training_on_client == 0:
                param.requires_grad = False
        elif name.startswith("bn1"):
            param.data = params_featrue_extractor[name]
            if args.whether_training_on_client == 0:
                param.requires_grad = False
        elif name.startswith("layer1"):
            param.data = params_featrue_extractor[name]
            if args.whether_training_on_client == 0:
                param.requires_grad = False
    # logging.info(client_model)
    return client_model

def create_server_model(n_classes):
    server_model = resnet56_server(n_classes)
    return server_model

def init_training_device(process_ID, fl_worker_num, gpu_num_per_machine):
    # initialize the mapping from process ID to GPU ID: <process ID, GPU ID>
    if process_ID == 0:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        return device
    process_gpu_dict = dict()
    for client_index in range(fl_worker_num):
        gpu_index = client_index % gpu_num_per_machine
        process_gpu_dict[client_index] = gpu_index

    # logging.info(process_gpu_dict)
    device = torch.device("cuda:" + str(process_gpu_dict[process_ID - 1]) if torch.cuda.is_available() else "cpu")
    logging.info(device)
    return device

if __name__ == "__main__":
    # initialize distributed computing (MPI)
    comm, process_id, worker_number = FedML_init()
    if process_id == 0:
        print("############################################0000000000000000####################")

    # parse python script input parameters
    parser = argparse.ArgumentParser()
    args = add_args(parser)
    # logging.info(args)

    # customize the process name
    str_process_name = "FedAvg (distributed):" + str(process_id)
    setproctitle.setproctitle(str_process_name)

    # customize the log format
    logging.basicConfig(level=logging.INFO,
                        format=str(
                            process_id) + ' - %(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                        datefmt='%a, %d %b %Y %H:%M:%S')
    hostname = socket.gethostname()
    logging.info("#############process ID = " + str(process_id) +
                 ", host name = " + hostname + "########" +
                 ", process ID = " + str(os.getpid()) +
                 ", process Name = " + str(psutil.Process(os.getpid())))

    # initialize the wandb machine learning experimental tracking platform (https://www.wandb.com/).
    if process_id == 0:
        wandb.init(
            project="knowledge-distillation",
            name="FedGKT-" + str(args.running_name),
            config=args,
            tags="group_knowledge_transfer"
        )

    # Set the random seed if provided (affects client sampling, and batching)
    # if a pseudorandom number generator is reinitialized with the same seed
    # it will produce the same sequence of numbers.
    seed = 0
    np.random.seed(seed)
    torch.manual_seed(np.random.randint(size))

    # GPU arrangement: Please customize this function according your own topology.
    # The GPU server list is configured at "mpi_host_file".
    # If we have 4 machines and each has two GPUs, and your FL network has 8 workers and a central worker.
    # The 4 machines will be assigned as follows:
    # machine 1: worker0, worker4, worker8;
    # machine 2: worker1, worker5;
    # machine 3: worker2, worker6;
    # machine 4: worker3, worker7;
    # Therefore, we can see that workers are assigned according to the order of machine list.
    logging.info("process_id = %d, size = %d" % (process_id, worker_number))
    device = init_training_device(process_id, worker_number - 1, args.gpu_num_per_server)

    # load data.
    # Note: if you use # of client epoch larger than 1,
    # please set the shuffle=False for the dataloader (CIFAR10/CIFAR100/CINIC10),
    # which keeps the batch sequence order across epoches.
    dataset = load_data(args, args.dataset)
    [train_data_num, test_data_num, train_data_global, test_data_global,
     train_data_local_num_dict, train_data_local_dict, test_data_local_dict, class_num] = dataset

    # create model
    if process_id == 0:
        model = create_server_model(class_num)
        logging.info("sever"+str(process_id) + "create_complete")
    else:
        model = create_client_model(args, class_num)
        logging.info("client" + str(process_id) + "create_complete")

    # start distributed training
    FedML_FedGKT_distributed(process_id, worker_number, device, comm, model, train_data_local_num_dict,
                             train_data_local_dict, test_data_local_dict, args)


import numpy as np
import tensorflow as tf
import os
import cv2

def get_images_labels_list():
    images_list = []
    labels_list = []
    with open("/home/linx-dl/hzr/2012_train.txt", 'r') as f:
        for line in f:
            image_full_name = line.split('\n')[0]
            # print(image_full_name)
            images_list.append(image_full_name)
            label_name = line.split('/')[-1].split('.')[0] + '.txt'
            label_full_name = os.path.join("/home/linx-dl/hzr/VOCdevkit/VOC2012/labels", label_name)

            #  print(label_full_name)
            labels_list.append(label_full_name)
    return images_list, labels_list

images_list,labels_list  = get_images_labels_list()

def parse_function(filename, label):
    image_string = tf.read_file(filename)

    label = tf.read_file(label)
    # label = tf.decode_raw(label,tf.uint8)
    # Don't use tf.image.decode_image, or the output shape will be undefined
    image = tf.image.decode_jpeg(image_string, channels=3)

    # This will convert to float values in [0, 1]
    image = tf.image.convert_image_dtype(image, tf.float32)

    image = tf.image.resize_images(image, [416, 416])
    return image, label

dataset = tf.data.Dataset.from_tensor_slices((tf.constant(images_list),tf.constant(labels_list)))
#dataset = dataset.map(lambda filename, label: tuple(tf.py_func(_read_py_function, [filename, label], [tf.uint8, tf.string])),num_parallel_calls=4)
dataset = dataset.map(parse_function,num_parallel_calls=4)
dataset = dataset.shuffle(len(images_list))
dataset = dataset.batch(10)
dataset = dataset.prefetch(1)

iterator = dataset.make_initializable_iterator()

with tf.Session() as sess:
    sess.run(iterator.initializer)

    batch_images, batch_labels = sess.run(iterator.get_next())

image_list=[]
for i,label in enumerate(batch_labels):
    label_list=[]
    for j,v in enumerate(label.decode().splitlines()):
        label_list.append([float(vv)for vv in v.split(' ')])
    image_list.append(label_list)

def iou_wh(r1, r2):
    min_w = min(r1[0], r2[0])
    min_h = min(r1[1], r2[1])
    area_r1 = r1[0] * r1[1]
    area_r2 = r2[0] * r2[1]

    intersect = min_w * min_h
    union = area_r1 + area_r2 - intersect

    return intersect / union

def get_active_anchors(roi, anchors):
    indxs = []
    iou_max, index_max = 0, 0

    for i, a in enumerate(anchors):
        iou = iou_wh(roi[1:], a)
        if iou > 0.7:
            indxs.append(i)
        if iou > iou_max:
            iou_max, index_max = iou, i

    if len(indxs) == 0:
        indxs.append(index_max)

    return indxs

def get_grid_cell(roi):
    x_center = roi[1] + roi[3] / 2.0
    y_center = roi[2] + roi[4] / 2.0

    grid_x = int(x_center / float(1 / 13))
    grid_y = int(y_center / float(1 / 13))

    return grid_x, grid_y

def roi2label(roi, anchor):
    x_center = roi[1] + roi[3] / 2.0
    y_center = roi[2] + roi[4] / 2.0

    grid_x = x_center / float(1 / 13)
    grid_y = y_center / float(1 / 13)

    grid_x_offset = grid_x - int(grid_x)
    grid_y_offset = grid_y - int(grid_y)

    roi_w_scale = roi[3] / (1 / anchor[0])
    roi_h_scale = roi[4] / (1 / anchor[1])

    label = [grid_x_offset, grid_y_offset, roi_w_scale, roi_h_scale]

    return label

label = np.zeros([10, 13, 13, 5, 6], dtype=np.float32)
anchors = [1.3221, 1.73145, 3.19275, 4.00944, 5.05587, 8.09892, 9.47112, 4.84053, 11.2364, 10.0071]
anchor_list = []
for i in range(0, len(anchors), 2):
    anchor_list.append([anchors[i] / 13, anchors[i + 1] / 13])
print(anchor_list)
for image_index in range(len(image_list)):  # every image
    for label_index in range(len(image_list[image_index])):  # every label
        # print(image_list[image_index][label_index])
        activate_anchor_index = get_active_anchors(image_list[image_index][label_index], anchor_list)

        grid_x, grid_y = get_grid_cell(image_list[image_index][label_index])
        for index in activate_anchor_index:
            anchor_label = roi2label(image_list[image_index][label_index], anchor_list[index])
            label[image_index, grid_x, grid_y, index] = np.concatenate(
                (anchor_label, [int(image_list[image_index][label_index][0])], [1.0]))

print(label[0,...])


# -*- coding: utf-8 -*-
# --------------------------------------
# @Time    : 2020/11/01
# @Author  : Oscar Chen
# @Email   : 530824679@qq.com
# @File    : train.py
# Description : train code
# --------------------------------------

import os
import math
import shutil
import numpy as np
import tensorflow as tf

from cfg.config import path_params, model_params, solver_params
from model.loss import Loss
from model.network import Network
from data import dataset, tfrecord
from utils.data_utils import *

def compute_warmup_lr(global_step, warmup_steps, name):
    """

    :param warmup_steps:
    :param name:
    :return:
    """
    with tf.variable_scope(name_or_scope=name):
        warmup_init_learning_rate = solver_params['init_learning_rate'] / 1000.0
        factor = tf.math.pow(solver_params['init_learning_rate'] / warmup_init_learning_rate, 1.0 / warmup_steps)
        warmup_lr = warmup_init_learning_rate * tf.math.pow(factor, global_step)
    return warmup_lr

def train():
    start_step = 0
    input_shape = model_params['input_shape']
    total_epoches = solver_params['total_epoches']
    batch_size = solver_params['batch_size']
    checkpoint_dir = path_params['checkpoints_dir']
    tfrecord_dir = path_params['tfrecord_dir']
    log_dir = path_params['logs_dir']
    initial_weight = path_params['initial_weight']
    restore = solver_params['restore']
    classes = read_class_names(path_params['class_file'])
    class_num = len(classes)

    # 创建相关目录
    ckpt_path = path_params['checkpoints_dir']
    if not os.path.exists(ckpt_path):
        os.makedirs(ckpt_path)

    logs_path = path_params['logs_dir']
    if os.path.exists(logs_path):
        shutil.rmtree(logs_path)
    os.makedirs(logs_path)

    # Set sess configuration
    gpu_options = tf.ConfigProto(allow_soft_placement=True)
    gpu_options.gpu_options.allow_growth = True
    gpu_options.gpu_options.allocator_type = 'BFC'

    # 解析得到训练样本以及标注
    data = tfrecord.TFRecord()
    train_tfrecord = os.path.join(tfrecord_dir, "train.tfrecord")
    data_num = total_sample(train_tfrecord)
    batch_num = int(math.ceil(float(data_num) / batch_size))
    dataset = data.create_dataset(train_tfrecord, batch_num, batch_size=batch_size, is_shuffle=False)

    # 创建训练和验证数据迭代器
    iterator = dataset.make_one_shot_iterator()
    inputs, y_true_20, y_true_40, y_true_80 = iterator.get_next()

    inputs.set_shape([None, input_shape[0], input_shape[1], 3])
    y_true_20.set_shape([None, 20, 20, 3, 5 + class_num])
    y_true_40.set_shape([None, 40, 40, 3, 5 + class_num])
    y_true_80.set_shape([None, 80, 80, 3, 5 + class_num])
    y_true = [y_true_20, y_true_40, y_true_80]

    # 构建网络计算损失
    network = Network(is_train=True)
    logits, preds = network.forward(inputs)

    losses = Loss()
    loss_op = losses.compute_loss(logits, preds, y_true, 'loss')

    #l2_loss = tf.losses.get_regularization_loss()
    #total_loss = loss_op[0] + loss_op[1] + loss_op[2]# + l2_loss

    # define training op
    global_step = tf.Variable(0, trainable=False)
    learning_rate = tf.train.exponential_decay(solver_params['init_learning_rate'], global_step, solver_params['decay_steps'], solver_params['decay_rate'], staircase=True)
    optimizer = tf.train.AdamOptimizer(learning_rate)
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = optimizer.minimize(loss_op[0], global_step=global_step)

    # 模型保存
    loader = tf.train.Saver(tf.moving_average_variables())
    save_variable = tf.global_variables()
    saver = tf.train.Saver(save_variable, max_to_keep=1000)

    # 配置tensorboard
    tf.summary.scalar('learn_rate', learning_rate)
    tf.summary.scalar("iou_loss", loss_op[1])
    tf.summary.scalar("conf_loss", loss_op[2])
    tf.summary.scalar("class_loss", loss_op[3])
    tf.summary.scalar('total_loss', loss_op[0])

    summary_op = tf.summary.merge_all()
    summary_writer = tf.summary.FileWriter(logs_path, graph=tf.get_default_graph(), flush_secs=60)

    # 开始训练
    with tf.Session(config=gpu_options) as sess:
        sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])

        if restore == True:
            ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
            if ckpt and ckpt.model_checkpoint_path:
                stem = os.path.basename(ckpt.model_checkpoint_path)
                restore_step = int(stem.split('.')[0].split('-')[-1])
                start_step = restore_step
                sess.run(global_step.assign(restore_step))
                loader.restore(sess, ckpt.model_checkpoint_path)
                print('Restoreing from {}'.format(ckpt.model_checkpoint_path))
            else:
                print("Failed to find a checkpoint")

        summary_writer.add_graph(sess.graph)
        try:
            print('=> Restoring weights from: %s ... ' % initial_weight)
            loader.restore(sess, initial_weight)
        except:
            print('=> %s does not exist !!!' % initial_weight)
            print('=> Now it starts to train from scratch ...')

        print('\n----------- start to train -----------\n')
        for epoch in range(start_step + 1, total_epoches):
            train_epoch_loss, train_epoch_iou_loss, train_epoch_confs_loss, train_epoch_class_loss = [], [], [], []
            for index in tqdm(range(batch_num)):
                _, summary_, loss_, iou_loss_, confs_loss_, class_loss_, global_step_, lr = sess.run(
                    [train_op, summary_op, loss_op[0], loss_op[1], loss_op[2], loss_op[3], global_step, learning_rate])

                train_epoch_loss.append(loss_)
                train_epoch_iou_loss.append(iou_loss_)
                train_epoch_confs_loss.append(confs_loss_)
                train_epoch_class_loss.append(class_loss_)

                summary_writer.add_summary(summary_, global_step_)

            train_epoch_loss, train_epoch_iou_loss, train_epoch_confs_loss, train_epoch_class_loss = np.mean(train_epoch_loss), np.mean(train_epoch_iou_loss), np.mean(train_epoch_confs_loss), np.mean(train_epoch_class_loss)
            print("Epoch: {}, global_step: {}, lr: {:.8f}, total_loss: {:.3f}, iou_loss: {:.3f},confs_loss: {:.3f}, class_loss: {:.3f}".format(epoch, global_step_, lr, train_epoch_loss, train_epoch_iou_loss, train_epoch_confs_loss, train_epoch_class_loss))
            snapshot_model_name = 'tusimple_train_miou={:.4f}.ckpt'.format(train_epoch_iou_loss)
            saver.save(sess, os.path.join(checkpoint_dir, snapshot_model_name), global_step=epoch)

        sess.close()

if __name__ == '__main__':
    train()


#!/usr/bin/env python

from __future__ import print_function
import sys
from threading import Lock

import chainer
import numpy as np

try:
    import chainer_mask_rcnn
except ImportError:
    print('''Please install chainer_mask_rcnn:

    sudo pip install chainer-mask-rcnn

''', file=sys.stderr)
    sys.exit(1)

import cv_bridge
from jsk_recognition_msgs.msg import ClusterPointIndices
from jsk_recognition_msgs.msg import Label
from jsk_recognition_msgs.msg import LabelArray
from jsk_topic_tools import ConnectionBasedTransport
from pcl_msgs.msg import PointIndices
import rospy
from sensor_msgs.msg import Image
from std_srvs.srv import Empty
from std_srvs.srv import EmptyResponse

class MaskRCNNInstanceSegmentation(ConnectionBasedTransport):

    def __init__(self):
        rospy.logwarn('This node is experimental, and its interface '
                      'can be changed in the future.')

        super(MaskRCNNInstanceSegmentation, self).__init__()
        # gpu
        self.gpu = rospy.get_param('~gpu', 0)
        if self.gpu >= 0:
            chainer.cuda.get_device_from_id(self.gpu).use()
        chainer.global_config.train = False
        chainer.global_config.enable_backprop = False

        self.fg_class_names = rospy.get_param('~fg_class_names')
        pretrained_model = rospy.get_param('~pretrained_model')

        n_fg_class = len(self.fg_class_names)
        self.model = chainer_mask_rcnn.models.MaskRCNNResNet(
            n_layers=50,
            n_fg_class=n_fg_class,
            pretrained_model=pretrained_model,
        )
        self.model.score_thresh = rospy.get_param('~score_thresh', 0.7)
        if self.gpu >= 0:
            self.model.to_gpu()

        self.lock = Lock()

        self.pub_indices = self.advertise(
            '~output/cluster_indices', ClusterPointIndices, queue_size=1)
        self.pub_labels = self.advertise(
            '~output/labels', LabelArray, queue_size=1)
        self.pub_lbl_cls = self.advertise(
            '~output/label_cls', Image, queue_size=1)
        self.pub_lbl_ins = self.advertise(
            '~output/label_ins', Image, queue_size=1)
        self.pub_viz = self.advertise(
            '~output/viz', Image, queue_size=1)

        self._bboxes = self._masks = self._labels = None
        self.srv_reset = rospy.Service('~reset', Empty, self.reset_callback)

    def subscribe(self):
        self.sub = rospy.Subscriber('~input', Image, self.callback,
                                    queue_size=1, buff_size=2**24)

    def unsubscribe(self):
        self.sub.unregister()

    def reset_callback(self, req):
        with self.lock:
            self._bboxes = None
            self._masks = None
            self._labels = None
        return EmptyResponse()

    def callback(self, imgmsg):
        bridge = cv_bridge.CvBridge()
        img = bridge.imgmsg_to_cv2(imgmsg, desired_encoding='rgb8')
        img_chw = img.transpose(2, 0, 1)  # C, H, W

        bboxes, masks, labels, scores = self.model.predict([img_chw])
        bboxes = bboxes[0]
        masks = masks[0]
        labels = labels[0]
        # scores = scores[0]

        self.lock.acquire()
        if self._bboxes is not None:
            assert self._masks is not None
            assert self._labels is not None
            assert len(self._bboxes) == len(self._masks) == len(self._labels)

            reuse = [True] * len(self._bboxes)
            for i, mask_i in enumerate(self._masks):
                for mask_j in masks:
                    iou = chainer_mask_rcnn.utils.get_mask_overlap(
                        mask_i, mask_j
                    )
                    if iou > 0.3:
                        reuse[i] = False
                        break
            bboxes = np.r_[bboxes, self._bboxes[reuse]]
            masks = np.r_[masks, self._masks[reuse]]
            labels = np.r_[labels, self._labels[reuse]]

        self._bboxes = bboxes
        self._masks = masks
        self._labels = labels
        self.lock.release()

        msg_indices = ClusterPointIndices(header=imgmsg.header)
        msg_labels = LabelArray(header=imgmsg.header)
        # -1: label for background
        lbl_cls = - np.ones(img.shape[:2], dtype=np.int32)
        lbl_ins = - np.ones(img.shape[:2], dtype=np.int32)
        for ins_id, (mask, label) in enumerate(zip(masks, labels)):
            indices = np.where(mask.flatten())[0]
            indices_msg = PointIndices(header=imgmsg.header, indices=indices)
            msg_indices.cluster_indices.append(indices_msg)
            class_name = self.fg_class_names[label]
            msg_labels.labels.append(Label(id=label, name=class_name))
            lbl_cls[mask] = label
            lbl_ins[mask] = ins_id  # instance_id
        self.pub_indices.publish(msg_indices)
        self.pub_labels.publish(msg_labels)

        msg_lbl_cls = bridge.cv2_to_imgmsg(lbl_cls)
        msg_lbl_ins = bridge.cv2_to_imgmsg(lbl_ins)
        msg_lbl_cls.header = msg_lbl_ins.header = imgmsg.header
        self.pub_lbl_cls.publish(msg_lbl_cls)
        self.pub_lbl_ins.publish(msg_lbl_ins)

        if self.pub_viz.get_num_connections() > 0:
            n_fg_class = len(self.fg_class_names)
            captions = ['{:d}: {:s}'.format(l, self.fg_class_names[l])
                        for l in labels]
            viz = chainer_mask_rcnn.utils.draw_instance_bboxes(
                img, bboxes, labels + 1, n_class=n_fg_class + 1,
                masks=masks, captions=captions)
            msg_viz = bridge.cv2_to_imgmsg(viz, encoding='rgb8')
            msg_viz.header = imgmsg.header
            self.pub_viz.publish(msg_viz)

if __name__ == '__main__':
    rospy.init_node('mask_rcnn_instance_segmentation')
    node = MaskRCNNInstanceSegmentation()
    rospy.spin()


import cv2 as cv
import sys
from datetime import datetime as dt
import time
import math
import asyncio
import queue
import os.path
import numpy as np
import imutils
from PySide2 import QtCore
from PySide2.QtCore import QThread, Signal
import ffmpeg
import logging
import subprocess

M0 = {
    "classesPath": "src/model/classes.names",
    "modelConfigurationPath": "src/model/yolov3-tiny.cfg",
    "modelWeightsPath": "src/model/yolov3-tiny.weights",
}

M1 = {
    "classesPath": "src/model/classes.names",
    "modelConfigurationPath": "src/model/yolov3-320.cfg",
    "modelWeightsPath": "src/model/yolov3.weights",
}

M2 = {
    "classesPath": "src/model/classes.names",
    "modelConfigurationPath": "src/model/yolov4.cfg",
    "modelWeightsPath": "src/model/yolov4.weights",
}

M3 = {
    "classesPath": "src/model/classes.names",
    "modelConfigurationPath": "src/model/yolov3-608.cfg",
    "modelWeightsPath": "src/model/yolov3.weights",
}

def printf(format, *args):
    sys.stdout.write(format % args)

class QueuingHandler(logging.Handler):
    def __init__(self, *args, message_queue, **kwargs):
        logging.Handler.__init__(self, *args, **kwargs)
        self.message_queue = message_queue

    def emit(self, record):
        self.message_queue.put(self.format(record).rstrip('\n'))

class VideoProcessor(QThread):

    on_data_finish = Signal(dict)
    on_progress = Signal(float)

    def __init__(self, sourceVideoPath, model):
        super().__init__()

        # Windows size in YOLO. Available options are: 320, 416, 608.
        # This field must be the same as in the YOLO*.cfg file.
        self.inpWidth = 416
        self.inpHeight = 416
        if model == 0:
            self.model = M0
            self.inpWidth = 416
            self.inpHeight = 416
            print("Chosen model: {}".format("YOLOv3-tiny"))
        elif model == 1:
            self.model = M1
            self.inpWidth = 320
            self.inpHeight = 320
            print("Chosen model: {}".format("YOLOv3-320"))
        elif model == 2:
            self.model = M2
            self.inpWidth = 416
            self.inpHeight = 416
            print("Chosen model: {}".format("YOLOv4 <Default>"))
        elif model == 3:
            self.model = M3
            self.inpWidth = 608
            self.inpHeight = 608
            print("Chosen model: {}".format("YOLOv3-608"))
        else:
            print("Wrong model. Can't find it!")

        self.confThreshold = 0.50   # Only objects with greater confidence will be detecting
        self.nmsThreshold = 0.7
        self.sourceVideoPath = sourceVideoPath
        # How many frames will be skipped.
        # E.g if value 3 is selected then only 1 frame per 3 frames will be analyzed.
        # 1 is a default value, no frames will be skipped.
        self.frameSkipRate = 1
        self.detectionsCache = []
        self.classes = []
        self.logEntries = []
        self.statistics = {
            "total_vehicles": 0,
            "light_vehicles": 0,
            "heavy_vehicles": 0,
            "two_wheel_vehicles": 0,
            "unknown_vehicles": 0
        }
        self.classesInterested = ['bicycle', 'motorbike', 'car', 'bus', 'truck']
        self.logger = logging.getLogger(__name__)

    def run(self):
        if not os.path.isfile(self.sourceVideoPath):
            return None

        classesName, modelConfiguration, modelWeights, *_ = self.model.values()
        classes = self.load_classes(classesName)
        self.classes = classes

        net = cv.dnn.readNetFromDarknet(modelConfiguration, modelWeights)
        net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)

        cap = self.get_video_capture(self.sourceVideoPath)
        self.fps = int(cap.get(cv.CAP_PROP_FPS))
        if self.fps < 10 and self.fps % 2 == 0:
            self.frameSkipRate = 2
        elif self.fps < 10 and self.fps % 2 != 0:
            self.frameSkipRate = 1
        elif self.fps > 10 and self.fps in [15,20,25,30,60]:
            self.frameSkipRate = 5
        else:
            self.frameSkipRate = 1
        self.current_frame = 0
        if os.name == 'nt':
            ffprobe_output = subprocess.check_output('powershell.exe "ffprobe -v error -count_frames -select_streams v:0 -show_entries stream=nb_read_frames -of default=nokey=1:noprint_wrappers=1 {0}"'.format(self.sourceVideoPath), stderr=subprocess.STDOUT, shell=True)
        elif os.name == 'posix':
            ffprobe_output = subprocess.check_output('ffprobe -v error -count_frames -select_streams v:0 -show_entries stream=nb_read_frames -of default=nokey=1:noprint_wrappers=1 {0}'.format(self.sourceVideoPath), stderr=subprocess.STDOUT, shell=True)
        ffprobe_output = ffprobe_output.rstrip().decode('utf-8')
        self.maximum = int(ffprobe_output)
        currFilename, currExtenstion = self.sourceVideoPath.split('.')
        writerVideoFile = f"{currFilename}_middle.{currExtenstion}"
        outputVideoFile = f"{currFilename}_rendered.{currExtenstion}"

        if (self.logger.hasHandlers()):
            self.logger.handlers.clear()

        logFormat = '%(asctime)s: %(name)8s: %(levelname)8s: %(message)s'
        outputLogFile = f"{currFilename}_rendered.log"
        logging.basicConfig(filename=outputLogFile, filemode='w', format=logFormat, level=logging.DEBUG)

        fH = logging.FileHandler(filename=outputLogFile, mode='w')

        message_queue = queue.Queue()
        qH = QueuingHandler(message_queue=message_queue, level=logging.DEBUG)

        formatter = logging.Formatter(logFormat)
        formatter.default_time_format = f"{dt.hour}:{dt.minute}:{dt.second}"
        qH.setFormatter(formatter)
        self.logger.addHandler(fH)
        self.logger.addHandler(qH)

        vid_writer = cv.VideoWriter(writerVideoFile,
                    cv.VideoWriter_fourcc(*'mp4v'), int(cap.get(cv.CAP_PROP_FPS) / self.frameSkipRate),
                    (int(cap.get(cv.CAP_PROP_FRAME_WIDTH)),
                    int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))))

        loop = asyncio.new_event_loop()
        try:
            if self.process_capture(cap, classes, net, vid_writer) == 0:
                self.logger.handlers.clear()
                self.log_detected_vehicles()
                logging.shutdown()
                vid_writer.release()
                interpolate = loop.create_task(self.interp_basic(writerVideoFile, outputVideoFile, self.fps))
                loop.run_until_complete(interpolate)
                loop.close()
                os.remove(writerVideoFile)
                self.on_progress.emit(self.current_frame / self.maximum)
                self.on_data_finish.emit({"done": True, "outVideo": outputVideoFile, "outLog": outputLogFile, "stats": self.statistics})
        except Exception as e:
            self.logger.error(str(e))
            self.logger.handlers.clear()
            logging.shutdown()
            self.on_data_finish.emit({"done": False, "outVideo": None, "outLog": None, "stats": None})

    def load_classes(self, path):
        # Load from coco.names classes specified in __init__
        if path.lower().endswith('.names') and path in self.model.values():
            with open(path, 'rt') as f:
                return f.read().rstrip('\n').split('\n')
        print("No classes loaded")
        return None

    def get_outputs_names(self, network):
        layersNames = network.getLayerNames()
        return [layersNames[i[0] - 1] for i in network.getUnconnectedOutLayers()]

    async def interp_basic(self, iVideo, oVideo, fps):
        input = ffmpeg.input(iVideo)
        video_interp = input.video.filter('minterpolate', fps)
        output = ffmpeg.output(video_interp, oVideo)
        ffmpeg.run(output, overwrite_output=True)

    def log_detected_vehicles(self):
        for element in self.logEntries:
            class_id = element[0]
            confidence = element[1]
            time = element[2]
            if self.classes[class_id] in self.classesInterested:
                if self.classes[class_id] == "car":
                    self.statistics["light_vehicles"] += 1
                elif self.classes[class_id] == "bus" or self.classes[class_id] == "truck":
                    self.statistics["heavy_vehicles"] += 1
                elif self.classes[class_id] == "bicycle" or self.classes[class_id] == "motorbike":
                    self.statistics["two_wheel_vehicles"] += 1
            self.statistics["total_vehicles"] += 1
            self.logger.info(f"New class: {self.classes[class_id]} detected in {time} second with conf: {confidence}")

    def create_log_entry(self, class_id, confidence, time):
        self.logEntries.append([class_id, confidence, time])

    def draw_pred(self, frame, classes, classId, conf, leftBtmX, leftBtmY, rightTopX, rightTopY):
        if classes[classId] not in self.classesInterested:
            return None

        # Draw a rectangle with lightblue line borders of thickness of 3 px
        cv.rectangle(frame, (leftBtmX, leftBtmY), (rightTopX, rightTopY), (255, 178, 50), 3)

        label = '%.2f' % conf

        if classes:
            label = '%s:%s' % (classes[classId], label) if classes[classId] in self.classesInterested else None

        confidence = label.split(':')[1]
        if [classId, confidence, self.current_frame/self.fps, leftBtmX, leftBtmY] not in self.logEntries:
            # Duplicate entries reduction
            self.create_log_entry(classId, confidence, self.current_frame/self.fps)

        labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        leftBtmY = max(leftBtmY, labelSize[1])
        # Draw a white label background
        cv.rectangle(frame, (leftBtmX, leftBtmY - round(1.5*labelSize[1])), (leftBtmX + round(1.5*labelSize[0]), leftBtmY + baseLine), (255, 255, 255), cv.FILLED)
        # Draw a black text label of a detected vehicle
        cv.putText(frame, label, (leftBtmX, leftBtmY), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 1)

    def post_process(self, frame, classes, outs):
        frameHeight = frame.shape[0]
        frameWidth = frame.shape[1]

        classIds = []
        confidences = []
        boxes = []

        # Parse coordinates of detected objects to rectangles indices.
        for out in outs:
            for detection in out:
                scores = detection[5:]
                classId = np.argmax(scores)
                confidence = scores[classId]
                if confidence > self.confThreshold:
                    # OpenCV returns relative indices of detected objects
                    # e.g if your frame is 1920px wide and center of the object is on 768px
                    # then OpenCV will return 0.4.
                    # Because every rectangle can be defined by two points on the surface
                    # we will obtain left bottom and width and height of a rectangle.
                    center_x = int(detection[0] * frameWidth)
                    center_y = int(detection[1] * frameHeight)
                    width = int(detection[2] * frameWidth)
                    height = int(detection[3] * frameHeight)
                    leftBtmX = int(center_x - width / 2)
                    leftBtmY = int(center_y - height / 2)
                    classIds.append(classId)
                    confidences.append(float(confidence))
                    boxes.append([leftBtmX, leftBtmY, width, height])

        indices = cv.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)
        for i in indices:
            i = i[0]
            box = boxes[i]
            leftBtmX = box[0]
            leftBtmY = box[1]
            width = box[2]
            height = box[3]
            # Draw rectangles on the frame
            self.draw_pred(frame, classes, classIds[i], confidences[i], leftBtmX, leftBtmY, leftBtmX + width, leftBtmY + height)

    def get_video_capture(self, path):
        if not os.path.isfile(path):
            print("Input file not found")
            sys.exit(1)
        else:
            return cv.VideoCapture(path)

    def process_capture(self, cap, classes, net, vid_writer):
        exit_status = 1
        start = time.time()
        while True:
            _, frame = cap.read()
            self.current_frame = cap.get(cv.CAP_PROP_POS_FRAMES)
            self.on_progress.emit(self.current_frame / self.maximum)
            frame = imutils.resize(frame)
            # Process frames in a grayscale.
            # Comment next 2 lines if you want to have a RGB colors
            frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
            frame = np.dstack([frame, frame, frame])

            # Processing has beed properly ended
            if int(cap.get(cv.CAP_PROP_POS_FRAMES)) == self.maximum:
                exit_status = 0
                break

            if self.current_frame % self.frameSkipRate:
                # vid_writer.write(frame.astype(np.uint8))
                continue

            blob = cv.dnn.blobFromImage(frame, 1/255, (self.inpWidth, self.inpHeight), [0,0,0], 1, crop=False)
            net.setInput(blob)
            outs = net.forward(self.get_outputs_names(net))
            self.post_process(frame, classes, outs)
            vid_writer.write(frame.astype(np.uint8))

        end = time.time()
        printf("Video Processing done in %.2f seconds.\n", (end-start))
        print(f"Total number of processed frames: {self.maximum}")
        printf("Mean frames per second rate: %.2f\n", self.maximum/(end-start))
        cap.release()
        cv.destroyAllWindows()
        return exit_status


# -*- coding: utf-8 -*-

"""Fit a potential to data with :mod:`~galpy`."""

__all__ = [
    "GalpyPotentialFitter",
    "GalpySCFPotentialFitter",
]

##############################################################################
# IMPORTS

# BUILT-IN
import typing as T

# THIRD PARTY
import astropy.coordinates as coord
import astropy.units as u
from galpy.potential import SCFPotential, scf_compute_coeffs_nbody

# PROJECT-SPECIFIC
import discO.type_hints as TH
from .wrapper import GalpyPotentialWrapper
from discO.core.fitter import PotentialFitter

# from discO.extern.galpy_potentials import scf_compute_coeffs_nbody
from discO.utils.coordinates import resolve_representationlike

##############################################################################
# PARAMETERS

GALPY_FITTER_REGISTRY: T.Dict[str, object] = dict()  # package : samplers

##############################################################################
# CODE
##############################################################################

class GalpyPotentialFitter(PotentialFitter, key="galpy"):
    """Fit a set of particles with Galpy.

    Parameters
    ----------
    potential_cls : object or str or None

    Other Parameters
    ----------------
    key : str or None

    """

    #######################################################
    # On the class

    _registry = GALPY_FITTER_REGISTRY

    #################################################################
    # On the instance

    def __new__(
        cls,
        *,
        potential_cls: T.Any = None,
        key: T.Optional[str] = None,
        **kwargs,
    ):
        # The class GalpyPotentialFitter is a wrapper for anything in its
        # registry If directly instantiating a GalpyPotentialFitter (not
        # subclass) we must also instantiate the appropriate subclass. Error
        # if can't find.
        if cls is GalpyPotentialFitter:

            # potential_cls overrides key
            key = potential_cls if isinstance(potential_cls, str) else key

            if key not in cls._registry:
                raise ValueError(
                    "PotentialFitter has no registered fitter for key: "
                    f"{key}",
                )

            # from registry. Registered in __init_subclass__
            kls = cls._registry[key]
            return kls.__new__(
                kls, potential_cls=potential_cls, key=None, **kwargs
            )

        elif key is not None:
            raise ValueError(f"Can't specify 'key' on {cls.__name__}.")

        return super().__new__(
            cls, potential_cls=potential_cls, key=None, **kwargs
        )

    # /def

    #######################################################
    # Fitting

    def __call__(
        self,
        c: TH.CoordinateType,
        mass: T.Optional[TH.QuantityType] = None,
        **kwargs,
    ) -> TH.SkyCoordType:
        """Fit Potential given particles.

        Parameters
        ----------
        c : coord-like

        frame: frame-like or None (optional, keyword-only)
            The frame of the fit potential.

            .. warning::

                Care should be taken that this matches the frame of the
                sampling potential.

        representation_type: |Representation| or None (optional, keyword-only)
            The coordinate representation.

        Returns
        -------
        :class:`~astropy.coordinates.SkyCoord`

        """
        raise NotImplementedError("Implement in subclass.")

    # /def

# /class

#####################################################################

class GalpySCFPotentialFitter(GalpyPotentialFitter, key="scf"):
    """Fit a set of particles with a Multipole expansion.

    Parameters
    ----------
    symmetry : str
    gridsizeR : int (optional)
    lmax : int (optional)
    **kwargs
        Passed to :class:`~galpy.potential.Potential`

    """

    def __new__(cls, **kwargs):
        kwargs.pop("potential_cls", None)
        kwargs.pop("key", None)
        return super().__new__(
            cls, potential_cls=SCFPotential, key=None, **kwargs
        )

    # /def

    def __init__(
        self,
        frame: TH.OptFrameLikeType = None,
        representation_type: T.Optional[TH.RepresentationType] = None,
        **kwargs,
    ) -> None:
        kwargs.pop("potential_cls", None)
        kwargs.pop("key", None)
        super().__init__(
            potential_cls=SCFPotential,
            key=None,
            frame=frame,
            representation_type=representation_type,
            **kwargs,
        )

    # /def

    #######################################################
    # Fitting

    def __call__(
        self,
        sample: TH.CoordinateType,
        mass: T.Optional[TH.QuantityType] = None,
        *,
        Nmax: int = None,
        Lmax: int = None,
        scale_factor: TH.QuantityType = None,
        **kwargs,
    ) -> TH.SkyCoordType:
        """Fit Potential given particles.

        .. todo::

            - amp != mass. Do this correctly.
            - work on ``n=[]`` array, not just niters

        Parameters
        ----------
        sample : |CoordinateFrame| or |SkyCoord|
        mass : |QuantityType|
        Nmax, Lmax : int or None (optional, keyword-only)
            The number of radial (N) and angular (L) coefficients.
            Must be integers > 0.
            If None (default) tries to draw from kwargs set at class
            initialization. If None set, raises ValueError.
        scale_factor : scalar |Quantity|
            units of distance or dimensionless

        Returns
        -------
        :class:`~astropy.coordinates.SkyCoord`

        Raises
        ------
        ValueError
            If `Nmax`, `Lmax` are None and no default set at initialization.

        """
        # --------------

        if mass is None:
            mass = sample.cache.get("mass")

        # kwargs
        kw = dict(self.potential_kwargs.items())  # deepcopy MappingProxyType
        kw.update(kwargs)

        # get from defaults if not passed
        _Nmax = kw.pop("Nmax", None)  # always try to pop
        Nmax = Nmax if Nmax is not None else _Nmax
        _Lmax = kw.pop("Lmax", None)  # always try to pop
        Lmax = Lmax if Lmax is not None else _Lmax

        if Nmax is None or Lmax is None:
            raise ValueError(
                "Nmax, Lmax have no default and must be specified.",
            )

        _scale_factor = kw.pop("scale_factor", 1 * u.one)
        scale_factor = (
            scale_factor if scale_factor is not None else _scale_factor
        )

        # --------------

        representation_type = resolve_representationlike(
            self.representation_type,
            error_if_not_type=False,
        )

        # --------------
        # Validation

        if Nmax < 0 or Lmax < 0:
            raise ValueError("Nmax & Lmax must be >=0.")

        if scale_factor.unit.physical_type not in ("length", "dimensionless"):
            raise u.UnitsError(
                "scale factor must have units of length or be dimensionless",
            )
        elif not scale_factor.isscalar:
            raise ValueError("scale factor must be a scalar.")

        # --------------

        sample = sample.transform_to(self.frame)
        position = sample.represent_as(coord.CartesianRepresentation).xyz

        # a dimensionless scale factor is assigned the same units as the
        # positions, so that (r / a) does not introduce an inadvertent scaling
        # from the units.
        if scale_factor.unit == u.one:
            scale_factor = scale_factor.value * position.unit

        # TODO don't do ``to_value`` when galpy supports units
        Acos, Asin = scf_compute_coeffs_nbody(
            position.to_value(position.unit),
            mass=mass.to_value(1e12 * u.solMass),
            N=Nmax,
            L=Lmax,
            a=scale_factor.to_value(position.unit),
            **kw,
        )

        return GalpyPotentialWrapper(
            self.potential_cls(
                amp=mass.sum(),
                Acos=Acos,
                Asin=Asin,
                a=scale_factor,
            ),
            frame=self.frame,
            representation_type=representation_type,
        )

    # /def

# /class

##############################################################################
# END