import numpy as np
import os
import tensorflow as tf

from .data_utils import minibatches, pad_sequences, get_chunks
from .general_utils import Progbar
from .base_model import BaseModel
from sklearn.metrics import precision_recall_fscore_support

class ASPECTModel(BaseModel):
    """Specialized class of Model for NER"""

    def __init__(self, config):
        super(ASPECTModel, self).__init__(config)
        self.idx_to_tag = {idx: tag for tag, idx in
                           self.config.vocab_tags.items()}
        self.max_length = 0

    def add_placeholders(self):
        """Define placeholders = entries to computational graph"""
        # shape = (batch size, max length of sentence in batch)
        self.word_ids = tf.placeholder(tf.int32, shape=[None, None],
                        name="word_ids")

        # shape = (batch size)
        self.sequence_lengths = tf.placeholder(tf.int32, shape=[None],
                        name="sequence_lengths")

        # shape = (batch size, max length of sentence, max length of word)
        self.char_ids = tf.placeholder(tf.int32, shape=[None, None, None],
                        name="char_ids")

        # shape = (batch_size, max_length of sentence)
        self.word_lengths = tf.placeholder(tf.int32, shape=[None, None],
                        name="word_lengths")

        # shape = (batch size, max length of sentence in batch)
        self.labels = tf.placeholder(tf.int32, shape=[None, None],
                        name="labels")

        # hyper parameters
        self.dropout = tf.placeholder(dtype=tf.float32, shape=[],
                        name="dropout")
        self.dropout_conv = tf.placeholder(dtype=tf.float32, shape=[],
                        name="dropout_conv")
        self.lr = tf.placeholder(dtype=tf.float32, shape=[],
                        name="lr")

    def get_feed_dict(self, words, labels=None, lr=None, dropout=None, dropout_conv=None):
        """Given some data, pad it and build a feed dictionary

        Args:
            words: list of sentences. A sentence is a list of ids of a list of
                words. A word is a list of ids
            labels: list of ids
            lr: (float) learning rate
            dropout: (float) keep prob

        Returns:
            dict {placeholder: value}

        """
        # perform padding of the given data
        if self.config.use_chars:
            char_ids, word_ids = zip(*words)
            word_ids, sequence_lengths = pad_sequences(word_ids, 0)
            char_ids, word_lengths = pad_sequences(char_ids, pad_tok=0,
                nlevels=2)
        else:
            word_ids, sequence_lengths,max_length = pad_sequences(words, 0)
            self.max_length = max_length

        # build feed dictionary
        feed = {
            self.word_ids: word_ids,
            self.sequence_lengths: sequence_lengths
        }

        if self.config.use_chars:
            feed[self.char_ids] = char_ids
            feed[self.word_lengths] = word_lengths

        if labels is not None:
            labels, _,_ = pad_sequences(labels, 0)
            feed[self.labels] = labels

        if lr is not None:
            feed[self.lr] = lr

        if dropout is not None:
            feed[self.dropout] = dropout

        if dropout_conv is not None:
            feed[self.dropout_conv] = dropout_conv

        return feed, sequence_lengths

    def add_word_embeddings_op(self):
        """Defines self.word_embeddings

        If self.config.embeddings is not None and is a np array initialized
        with pre-trained word vectors, the word embeddings is just a look-up
        and we don't train the vectors. Otherwise, a random matrix with
        the correct shape is initialized.
        """
        with tf.variable_scope("words"):
            if self.config.embeddings is None:
                self.logger.info("WARNING: randomly initializing word vectors")
                _word_embeddings = tf.get_variable(
                        name="_word_embeddings",
                        dtype=tf.float32,
                        shape=[self.config.nwords, self.config.dim_word])
            else:
                _word_embeddings = tf.Variable(
                        self.config.embeddings,
                        name="_word_embeddings",
                        dtype=tf.float32,
                        trainable=self.config.train_embeddings)

            word_embeddings = tf.nn.embedding_lookup(_word_embeddings,
                    self.word_ids, name="word_embeddings")

        with tf.variable_scope("chars"):
            if self.config.use_chars:
                # get char embeddings matrix
                _char_embeddings = tf.get_variable(
                        name="_char_embeddings",
                        dtype=tf.float32,
                        shape=[self.config.nchars, self.config.dim_char])
                char_embeddings = tf.nn.embedding_lookup(_char_embeddings,
                        self.char_ids, name="char_embeddings")

                # put the time dimension on axis=1
                s = tf.shape(char_embeddings)
                char_embeddings = tf.reshape(char_embeddings,
                        shape=[s[0]*s[1], s[-2], self.config.dim_char])
                word_lengths = tf.reshape(self.word_lengths, shape=[s[0]*s[1]])

                # bi lstm on chars
                cell_fw = tf.contrib.rnn.LSTMCell(self.config.hidden_size_char,
                        state_is_tuple=True)
                cell_bw = tf.contrib.rnn.LSTMCell(self.config.hidden_size_char,
                        state_is_tuple=True)
                _output = tf.nn.bidirectional_dynamic_rnn(
                        cell_fw, cell_bw, char_embeddings,
                        sequence_length=word_lengths, dtype=tf.float32)

                # read and concat output
                _, ((_, output_fw), (_, output_bw)) = _output
                output = tf.concat([output_fw, output_bw], axis=-1)

                # shape = (batch size, max sentence length, char hidden size)
                output = tf.reshape(output,
                        shape=[s[0], s[1], 2*self.config.hidden_size_char])
                word_embeddings = tf.concat([word_embeddings, output], axis=-1)

        self.word_embeddings =  tf.nn.dropout(word_embeddings, self.dropout)

        #print(tf.Print(self.word_embeddings,[self.word_embeddings]))

        if self.config.conv:
            self.temp = tf.squeeze(tf.extract_image_patches(self.word_embeddings[:,:,:,tf.newaxis], ksizes=[1, self.config.WINDOW_LEN, self.config.DIM, 1], strides=[1, self.config.stride, self.config.DIM, 1], rates=[1, 1, 1, 1], padding='SAME'))
            self.image_patches = tf.reshape(self.temp, (-1, tf.shape(self.word_ids)[1], self.config.WINDOW_LEN, self.config.DIM))
            self.image_patches_reshaped = tf.reshape(self.image_patches, (-1, self.config.WINDOW_LEN, self.config.DIM))[:,:,:,tf.newaxis]

    def add_logits_op(self):
        """Defines self.logits

        For each word in each sentence of the batch, it corresponds to a vector
        of scores, of dimension equal to the number of tags.
        """
        with tf.variable_scope("bi-lstm"):
            cell_fw = tf.contrib.rnn.LSTMCell(self.config.hidden_size_lstm)
            cell_bw = tf.contrib.rnn.LSTMCell(self.config.hidden_size_lstm)
            (output_fw, output_bw), _ = tf.nn.bidirectional_dynamic_rnn(
                    cell_fw, cell_bw, self.word_embeddings,
                    sequence_length=self.sequence_lengths, dtype=tf.float32)
            output = tf.concat([output_fw, output_bw], axis=-1)
            output = tf.nn.dropout(output, self.dropout)

        with tf.variable_scope("proj"):
            W = tf.get_variable("W", dtype=tf.float32,
                    shape=[2*self.config.hidden_size_lstm, self.config.ntags])

            b = tf.get_variable("b", shape=[self.config.ntags],
                    dtype=tf.float32, initializer=tf.zeros_initializer())

            nsteps = tf.shape(output)[1]
            output = tf.reshape(output, [-1, 2*self.config.hidden_size_lstm])
            pred = tf.matmul(output, W) + b
            self.logits = tf.reshape(pred, [-1, nsteps, self.config.ntags])
    def add_mlp_op(self):
        """Defines self.logits

        For each word in each sentence of the batch, it corresponds to a vector
        of scores, of dimension equal to the number of tags.
        """
        with tf.variable_scope("mlp"):
            dense_input = tf.reshape(self.word_embeddings, (-1, self.config.dim_word))
            #dense_input = tf.nn.dropout(dense_input, self.dropout_conv)
            output = tf.contrib.layers.fully_connected(
            dense_input,
            self.config.mlp_size,
            activation_fn=tf.nn.relu,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            #weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input3"
            )
            output = tf.contrib.layers.fully_connected(
            output,
            self.config.ntags,
            activation_fn=None,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input4"
            )
        self.logits = tf.reshape(output, (-1, tf.shape(self.word_ids)[1], self.config.ntags))
    def add_mlp_op_conv(self):
        """Defines self.logits

        For each word in each sentence of the batch, it corresponds to a vector
        of scores, of dimension equal to the number of tags.
        """
        with tf.variable_scope("mlp"):
            dense_input = tf.reshape(self.image_patches_reshaped, (-1, self.config.WINDOW_LEN * self.config.dim_word))
            #dense_input = tf.nn.dropout(dense_input, self.dropout_conv)
            output = tf.contrib.layers.fully_connected(
            dense_input,
            self.config.mlp_size,
            activation_fn=tf.nn.relu,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            #weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input1"
            )
            output = tf.nn.dropout(output, self.dropout)
            output = tf.contrib.layers.fully_connected(
            output,
            self.config.ntags,
            activation_fn=None,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input2"
            )
        self.logits = tf.reshape(output, (-1, tf.shape(self.word_ids)[1], self.config.ntags))

    def add_logits_op_conv(self):
        """Defines self.logits

        For each word in each sentence of the batch, it corresponds to a vector
        of scores, of dimension equal to the number of tags.
        """

        with tf.name_scope("conv-maxpool"):

                 # ( (BATCH_SIZE*WORDS), WINDOW_LEN, DIM, 1 )
            pooled_out = []
            for i,filter_size in enumerate(self.config.FILTER_SIZE):
                filter_shape = [filter_size, self.config.DIM, 1, self.config.NUMBER_OF_FEATURE_MAPS[i]]
                W = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W")
                    #W = tf.get_variable(shape = filter_shape, initializer = tf.truncated_normal_initializer(stddev=0.001), name="W"+str(i))
                b = tf.Variable(tf.constant(0.1, shape=[self.config.NUMBER_OF_FEATURE_MAPS[i]]), name="b")

                conv = tf.nn.conv2d(self.image_patches_reshaped,filter=W,strides=[1, 1, 1, 1],padding="VALID",name="conv")
                    #print(tf.Print(conv,[conv]))
    # conv = tf.squeeze(conv) # ( (BATCH_SIZE*WORDS), WINDOW_LEN-FILTER_SIZE + 1, NUMBER_OF_FEATURE_MAPS)
                    #conv = tf.nn.bias_add(conv,b)
                    #conv = tf.nn.relu(conv)
                    #conv_non_linear = tf.nn.relu(tf.nn.bias_add(conv, b), name="relu") # ( (BATCH_SIZE*WORDS), WINDOW_LEN-FILTER_SIZE + 1, 1, NUMBER_OF_FEATURE_MAPS)

                pooled = tf.nn.max_pool(conv,ksize=[1, (self.config.WINDOW_LEN-filter_size+1), 1, 1],strides=[1, 1, 1, 1],padding='VALID',data_format = 'NHWC',name="pool")
                pooled = tf.squeeze(pooled) # ( (BATCH_SIZE*WORDS), NUMBER_OF_FEATURE_MAPS)
                self.output = tf.reshape(pooled, (-1, tf.shape(self.word_ids)[1], self.config.NUMBER_OF_FEATURE_MAPS[i]))
                pooled_out.append(self.output)

        self.h_pool = tf.concat(pooled_out, 2)

        with tf.name_scope("size_calc"):
            size=0
            for i in range(len(self.config.FILTER_SIZE)):
                size += self.config.NUMBER_OF_FEATURE_MAPS[i]

        with tf.name_scope("conv2-maxpool"):
            filter_shape = [self.config.conv2_filter_size, size, self.config.conv2_dim]
            W = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W_")
            b = tf.Variable(tf.constant(0.1, shape=[self.config.conv2_dim]), name="b_")
            conv_ = tf.nn.conv1d(self.h_pool,filters=W,stride=1,padding="SAME",name="conv_")
            #conv_ = tf.squeeze(conv_)

        with tf.variable_scope("proj"):

            dense_input = tf.reshape(self.h_pool, (-1, size))
            #dense_input = tf.nn.dropout(dense_input, self.dropout_conv)

            output = tf.contrib.layers.fully_connected(
            dense_input,
            self.config.mlp_size,
            activation_fn=tf.nn.relu,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input1"
            )

            # dense_input = tf.reshape(self.image_patches_reshaped, (-1, self.config.WINDOW_LEN * self.config.dim_word))
            # #dense_input = tf.nn.dropout(dense_input, self.dropout_conv)
            # output2 = tf.contrib.layers.fully_connected(
            # dense_input,
            # self.config.mlp_size,
            # activation_fn=tf.nn.relu,
            # normalizer_fn=None,
            # normalizer_params=None,
            # weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            # #weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            # biases_initializer=tf.zeros_initializer(),
            # trainable=True,
            # scope="input3"
            # )

            # #output = tf.concat([tf.reshape(output,(-1, tf.shape(self.word_ids)[1], self.config.mlp_size)),tf.reshape(output2,(-1, tf.shape(self.word_ids)[1], self.config.mlp_size))],axis=2)

            # #output = tf.reshape(output, (-1, tf.shape(self.word_ids)[1]*2*self.config.mlp_size))

            # output = tf.concat([output, output2],axis = 1)

            output = tf.nn.dropout(output, self.dropout_conv)

            output = tf.contrib.layers.fully_connected(
            output,
            self.config.ntags,
            activation_fn=None,
            normalizer_fn=None,
            normalizer_params=None,
            weights_initializer=tf.contrib.layers.xavier_initializer(uniform=True, seed=1227),
            weights_regularizer=tf.contrib.layers.l2_regularizer(0.001),
            biases_initializer=tf.zeros_initializer(),
            trainable=True,
            scope="input2"
            )

        self.logits = tf.reshape(output, (-1, tf.shape(self.word_ids)[1], self.config.ntags))

    def add_pred_op(self):
        """Defines self.labels_pred

        This op is defined only in the case where we don't use a CRF since in
        that case we can make the prediction "in the graph" (thanks to tf
        functions in other words). With theCRF, as the inference is coded
        in python and not in pure tensroflow, we have to make the prediciton
        outside the graph.
        """
        if not self.config.use_crf:
            self.labels_pred = tf.cast(tf.argmax(self.logits, axis=-1),
                    tf.int32)

    def add_loss_op(self):
        """Defines the loss"""
        if self.config.use_crf:
            log_likelihood, trans_params = tf.contrib.crf.crf_log_likelihood(
                    self.logits, self.labels, self.sequence_lengths)
            self.trans_params = trans_params # need to evaluate it for decoding
            self.loss = tf.reduce_mean(-log_likelihood)
        else:
            losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
                    logits=self.logits, labels=self.labels)
            mask = tf.sequence_mask(self.sequence_lengths)
            losses = tf.boolean_mask(losses, mask)
            self.loss = tf.reduce_mean(losses)

        # for tensorboard
        tf.summary.scalar("loss", self.loss)

    def build(self):
        # NER specific functions
        self.add_placeholders()
        self.add_word_embeddings_op()
        #self.add_logits_op()
        self.add_logits_op_conv()
        #self.add_mlp_op()
        #self.add_mlp_op_conv()
        self.add_pred_op()
        self.add_loss_op()

        # Generic functions that add training op and initialize session
        self.add_train_op(self.config.lr_method, self.lr, self.loss,
                self.config.clip)
        self.initialize_session() # now self.sess is defined and vars are init

    def predict_batch(self, words):
        """
        Args:
            words: list of sentences

        Returns:
            labels_pred: list of labels for each sentence
            sequence_length

        """
        fd, sequence_lengths = self.get_feed_dict(words, dropout=1.0,dropout_conv=1.0)

        if self.config.use_crf:
            # get tag scores and transition params of CRF
            viterbi_sequences = []
            logits, trans_params = self.sess.run(
                    [self.logits, self.trans_params], feed_dict=fd)

            # iterate over the sentences because no batching in vitervi_decode
            for logit, sequence_length in zip(logits, sequence_lengths):
                logit = logit[:sequence_length] # keep only the valid steps
                viterbi_seq, viterbi_score = tf.contrib.crf.viterbi_decode(
                        logit, trans_params)
                viterbi_sequences += [viterbi_seq]

            return viterbi_sequences, sequence_lengths

        else:
            labels_pred = self.sess.run(self.labels_pred, feed_dict=fd)

            return labels_pred, sequence_lengths

    def run_epoch(self, train, dev, epoch):
        """Performs one complete pass over the train set and evaluate on dev

        Args:
            train: dataset that yields tuple of sentences, tags
            dev: dataset
            epoch: (int) index of the current epoch

        Returns:
            f1: (python float), score to select model on, higher is better

        """
        # progbar stuff for logging
        batch_size = self.config.batch_size
        nbatches = (len(train) + batch_size - 1) // batch_size
        prog = Progbar(target=nbatches)

        # iterate over dataset
        for i, (words, labels) in enumerate(minibatches(train, batch_size)):
            fd, _ = self.get_feed_dict(words, labels, self.config.lr,
                    self.config.dropout, self.config.dropout_conv)

            _, train_loss, summary = self.sess.run(
                    [self.train_op, self.loss, self.merged], feed_dict=fd)

            prog.update(i + 1, [("train loss", train_loss)])

            # tensorboard
            if i % 10 == 0:
                self.file_writer.add_summary(summary, epoch*nbatches + i)

        metrics = self.run_evaluate(dev)
        msg = " - ".join(["{} {}".format(k, v)
                for k, v in metrics.items()])
        self.logger.info(msg)

        return metrics["f1"]

    def run_evaluate(self, test):
        """Evaluates performance on test set

        Args:
            test: dataset that yields tuple of (sentences, tags)

        Returns:
            metrics: (dict) metrics["acc"] = 98.4, ...

        """
        accs = []
        correct_preds, total_correct, total_preds = 0., 0., 0.
        gold = []
        pred = []
        for words, labels in minibatches(test, self.config.batch_size):
            labels_pred, sequence_lengths = self.predict_batch(words)

            for lab, lab_pred, length in zip(labels, labels_pred,
                                             sequence_lengths):
                lab      = lab[:length]
                lab_pred = lab_pred[:length]
                gold+=lab
                pred+=lab_pred
                accs    += [a==b for (a, b) in zip(lab, lab_pred)]
                lab_chunks      = set(get_chunks(lab, self.config.vocab_tags,message = "gold standard"))
                lab_pred_chunks = set(get_chunks(lab_pred,
                                                 self.config.vocab_tags,message = "prediction"))

                correct_preds += len(lab_chunks & lab_pred_chunks)
                total_preds   += len(lab_pred_chunks)
                total_correct += len(lab_chunks)

        p   = correct_preds / total_preds if correct_preds > 0 else 0
        r   = correct_preds / total_correct if correct_preds > 0 else 0
        f1  = 2 * p * r / (p + r) if correct_preds > 0 else 0
        acc = np.mean(accs)
        score = precision_recall_fscore_support(gold, pred, average='macro')
        #score_bi = precision_recall_fscore_support(gold, pred, labels =[1,2], average='macro')
        return {"acc": 100*acc, "f1": 100*f1, "partial_matching_stat": score}

    def predict(self, words_raw):
        """Returns list of tags

        Args:
            words_raw: list of words (string), just one sentence (no batch)

        Returns:
            preds: list of tags (string), one for each word in the sentence

        """
        words = [self.config.processing_word(w) for w in words_raw]
        if type(words[0]) == tuple:
            words = zip(*words)
        pred_ids, _ = self.predict_batch([words])
        preds = [self.idx_to_tag[idx] for idx in list(pred_ids[0])]

        return preds


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

import torch

from utils import config
from numpy import random
from models.layers import Encoder
from models.layers import Decoder
from models.layers import ReduceState
from transformer.model import TranEncoder

use_cuda = config.use_gpu and torch.cuda.is_available()

random.seed(123)
torch.manual_seed(123)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(123)

class Model(object):
    def __init__(self, model_path=None, is_eval=False, is_tran = False):
        encoder = Encoder()
        decoder = Decoder()
        reduce_state = ReduceState()
        if is_tran:
            encoder = TranEncoder(config.vocab_size, config.max_enc_steps, config.emb_dim,
                 config.n_layers, config.n_head, config.d_k, config.d_v, config.d_model, config.d_inner)

        # shared the embedding between encoder and decoder
        decoder.tgt_word_emb.weight = encoder.src_word_emb.weight

        if is_eval:
            encoder = encoder.eval()
            decoder = decoder.eval()
            reduce_state = reduce_state.eval()

        if use_cuda:
            encoder = encoder.cuda()
            decoder = decoder.cuda()
            reduce_state = reduce_state.cuda()

        self.encoder = encoder
        self.decoder = decoder
        self.reduce_state = reduce_state

        if model_path is not None:
            state = torch.load(model_path, map_location=lambda storage, location: storage)
            self.encoder.load_state_dict(state['encoder_state_dict'])
            self.decoder.load_state_dict(state['decoder_state_dict'], strict=False)
            self.reduce_state.load_state_dict(state['reduce_state_dict'])


# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import numpy as np
import paddle
import paddle.fluid as fluid
import paddle.fluid.layers as FL
from paddle.fluid.param_attr import ParamAttr

from . import resnet
from . import lesionnet
from . import densenet
from . import inception_v4 as inception

from learning_rate import cosine_decay_with_warmup

__all__ = ["create_model"]

ResNetModels = sorted(name for name in resnet.__dict__
                             if callable(resnet.__dict__[name]))
DenseNetModels = sorted(name for name in densenet.__dict__
                             if callable(densenet.__dict__[name]))

class create_model():
    def __init__(self, data_shape, loss_type, main_arch, name, model_type='train', **kwargs):
        """初始化函数"""
        self.loss_type = loss_type
        self.data_shape = data_shape
        self.main_arch = main_arch
        self.name = name
        self.model_type = model_type

    def create_acc_op(self, predict, label):
        """compute accuracy with tensor

         Args:
         predict: model output tensor activated by softmax
         label: a non-sparse tensor

         Returns:
         acc: acc tensor
         """
        accuracy = FL.accuracy(input=predict, label=label)
        return accuracy

    def create_reader_op(self, feed_list):
        """create reader with op

         Args:
         feed_list: placeholders needs to be feed

         Returns:
         reader: a reader
         """
        reader = fluid.io.DataLoader.from_generator(feed_list=feed_list,
                     capacity=256, iterable=True, use_double_buffer=True)

        return reader

    def create_loss_op(self, predict, label, epsilon=1e-7):
        """compute loss with tensor

         Args:
         predict: model output tensor activated by softmax
         label: a non-sparse tensor

         Returns:
         loss: cross-entropy loss
         """
        if self.loss_type == "nl" and self.model_type == "train":
            one_hot_label = fluid.one_hot(label, depth=predict.shape[-1])
            one_hot_label = FL.squeeze(one_hot_label, axes=[-2])
            # log
            neg_prob = 1 - predict
            log_neg_prob = FL.log(fluid.layers.clip(neg_prob, min=epsilon, max=1.))
            ce_loss = -1 * log_neg_prob * one_hot_label
            cost = FL.reduce_sum(ce_loss, dim=-1, keep_dim=True)
        else: # PL or evaluation
            cost = FL.cross_entropy(predict, label)

        loss = FL.mean(cost)

        return loss

    def create_cam_op(self, predict, class_dim, heatmaps):
        """compute loss with tensor

         Args:
         predict: model output tensor activated by softmax
         class_dim: dim of multi-class vector
         heatmaps: 全局池化前的特征图

         Returns:
         heatmaps: class activation map
         """
        if self.main_arch in DenseNetModels:
            weights_shape = 1024
            name = "fc_weights"
        elif self.main_arch == "xception":
            weights_shape = 2048
            name = "fc_weights"
        else:
            raise ValueError("Calc CAM of model arch {} is not supported.".format(self.main_arch))

        fc_weights = FL.create_parameter(shape=[weights_shape, class_dim], dtype='float32', name=name) # 1024, 5

        pred_idx = FL.argmax(predict, 1) # bs, 1
        fc_weights = FL.transpose(fc_weights, perm=[1, 0]) # 5, 1024
        fc_weights = FL.gather(fc_weights, index=pred_idx) # bs, 1024

        heatmaps = heatmaps * fc_weights # bs, 1024, 16, 16
        heatmaps = FL.reduce_sum(heatmaps, dim=1, keep_dim=False)

        return heatmaps

    def net(self, class_dim=5, CAM=False):
        """Create second stage model
         Args:
         class_dim: dim of multi-class vector
         CAM:  是否创建CAM heatmap
         Returns:
         * A list contain 4/5 tensors / ops:
             - loss, cross-entropy loss tensor
             - accuracy, accuracy metric tensor
             - predict, model output tensor activated by softmax
             - hacked_img_id, img_id tensor
             - cam_heatmap, only if CAM == True, class activation map tensor
         * reader, reader op to feed data into placeholder
         """
        self.input_feature = fluid.data(name='{}_input'.format(self.name), shape=[-1] + self.data_shape, dtype='uint8')
        self.label = fluid.data(name='{}_label'.format(self.name), shape=[-1, 1], dtype='int64')
        self.img_id = fluid.data(name='{}_img_id'.format(self.name), shape=[-1, 1], dtype='int64')

        # Lesion Net
        lesion = lesionnet.LesionNet()

        # Backbone
        if self.main_arch in ResNetModels:
            model = resnet.__dict__[self.main_arch]()
        elif self.main_arch in DenseNetModels:
            model = densenet.__dict__[self.main_arch]()
        elif self.main_arch == "inception":
            model = inception.InceptionV4()
        else:
            raise ValueError("Model {} is not supported.".format(self.main_arch))

        inp = FL.transpose(FL.cast(self.input_feature, "float32"), perm=[0, 3, 1, 2]) / 255.

        # Element wise mul of lesion prob maps and input image
        lesion_probs = lesion.net(inp, class_dim=4) # bs, 4, 16, 16
        lesion_probs = FL.split(lesion_probs, num_or_sections=4, dim=1) # probs, bs*1*16*16 4

        I = FL.image_resize(inp, out_shape=(512, 512), resample="BILINEAR")
        Is = []

        for L in lesion_probs:
            W = FL.image_resize(L, out_shape=(512, 512), resample="NEAREST") # bs, 1, 512, 512
            temp_I = FL.elementwise_mul(I, FL.expand(W + 1., expand_times=[1, 3, 1, 1])) # W + 1., bs, 3, 512, 512
            Is.append(temp_I)
        I = FL.concat(Is, axis=1) # bs, 3*4, 512, 512
        I.stop_gradient = True

        lesion_pos_prob = 1. - lesion_probs[0]
        main_arch_out = model.net(I, class_dim=class_dim, lesion_map=lesion_pos_prob, CAM=CAM)

        if CAM:
            logit, heatmaps = main_arch_out
        else:
            logit = main_arch_out

        predict = FL.softmax(logit)
        accuracy = self.create_acc_op(predict, self.label)
        loss = self.create_loss_op(predict, self.label)
        reader = self.create_reader_op([self.img_id, self.input_feature, self.label])

        # This is a hack
        hacked_img_id = FL.cast(self.img_id, "int32")

        if CAM:
            cam_heatmap = self.create_cam_op(predict, class_dim, heatmaps)
            return [loss, accuracy, predict, hacked_img_id, cam_heatmap], reader

        return [loss, accuracy, predict, hacked_img_id], reader


import json
import argparse
import torch
import os
import random
import numpy as np
import requests
import logging
import math
import copy
import string
import faiss

from time import time
from tqdm import tqdm, trange
from densephrases.utils.squad_utils import get_question_dataloader
from densephrases.utils.single_utils import load_encoder
from densephrases.utils.open_utils import load_phrase_index, get_query2vec, load_qa_pairs
from densephrases.utils.eval_utils import drqa_exact_match_score, drqa_regex_match_score, \
                                          drqa_metric_max_over_ground_truths
from eval_phrase_retrieval import evaluate, evaluate_type
from densephrases import Options

from transformers import (
    AdamW,
    get_linear_schedule_with_warmup,
)

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s', datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

global_acc = []

def train_query_encoder(args, mips=None):
    # Freeze one for MIPS
    device = 'cuda' if args.cuda else 'cpu'
    logger.info("Loading pretrained encoder: this one is for MIPS (fixed)")
    pretrained_encoder, tokenizer, _ = load_encoder(device, args)
    pretrained_encoder.eval()

    # Train a copy of it
    logger.info("Copying target encoder")
    target_encoder = copy.deepcopy(pretrained_encoder)

    # MIPS
    if mips is None:
        mips = load_phrase_index(args)

    # Optimizer setting
    # layers_to_freeze = [f"layer.{i}." for i in range(0, args.num_layers_freeze)]
    layers_to_freeze = []
    def is_train_param(name):
        if name.endswith(".embeddings.word_embeddings.weight"):
            logger.info(f'freezing {name}')
            return False

        if any(nd in name for nd in layers_to_freeze):
            logger.info(f'freezing {name}')
            return False
        return True

    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [{
            "params": [
                p for n, p in target_encoder.named_parameters() \
                    if not any(nd in n for nd in no_decay) and is_train_param(n)
            ],
            "weight_decay": args.weight_decay,
        }, {
            "params": [
                p for n, p in target_encoder.named_parameters() \
                    if any(nd in n for nd in no_decay) and is_train_param(n)
            ],
            "weight_decay": 0.0
        },
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    step_per_epoch = math.ceil(len(load_qa_pairs(args.train_path, args)[1]) / args.per_gpu_train_batch_size)
    t_total = int(step_per_epoch // args.gradient_accumulation_steps * args.num_train_epochs)
    logger.info(f"Train for {t_total} iterations")
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
     )
    eval_steps = math.ceil(len(load_qa_pairs(args.dev_path, args)[1]) / args.eval_batch_size)
    logger.info(f"Test takes {eval_steps} iterations")

    #qids, questions, answers, _ = load_type_qa_pairs(args.test_path, args, q_types=args.question_types, q_idx=q_idx)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        target_encoder, optimizer = amp.initialize(target_encoder, optimizer, opt_level=args.fp16_opt_level)

    # Train arguments
    best_acc = -1000.0
    global_acc.append([0, 0, best_acc])
    for ep_idx in range(int(args.num_train_epochs)):

        # Training
        total_loss = 0.0
        total_accs = []
        total_accs_k = []

        # Load training dataset
        q_ids, questions, answers, titles = load_qa_pairs(args.train_path, args, shuffle=True)
        pbar = tqdm(get_top_phrases(
            mips, q_ids, questions, answers, titles, target_encoder, tokenizer,
            args.per_gpu_train_batch_size, args)
        )

        for step_idx, (q_ids, questions, answers, titles, outs) in enumerate(pbar):
            train_dataloader, _, _ = get_question_dataloader(
                questions, tokenizer, args.max_query_length, batch_size=args.per_gpu_train_batch_size
            )
            svs, evs, tgts, p_tgts = annotate_phrase_vecs(mips, q_ids, questions, answers, titles, outs, args)

            target_encoder.train()
            svs_t = torch.Tensor(svs).to(device)
            evs_t = torch.Tensor(evs).to(device)
            tgts_t = [torch.Tensor([tgt_ for tgt_ in tgt if tgt_ is not None]).to(device) for tgt in tgts]
            p_tgts_t = [torch.Tensor([tgt_ for tgt_ in tgt if tgt_ is not None]).to(device) for tgt in p_tgts]
            logger.info(f"tgts_t {tgts_t}")
            logger.info(f"p_tgts_t {p_tgts_t}")

            # Train query encoder
            assert len(train_dataloader) == 1
            for batch in train_dataloader:
                batch = tuple(t.to(device) for t in batch)
                loss, accs = target_encoder.train_query(
                    input_ids_=batch[0], attention_mask_=batch[1], token_type_ids_=batch[2],
                    start_vecs=svs_t,
                    end_vecs=evs_t,
                    targets=tgts_t,
                    p_targets=p_tgts_t,
                )

                # Optimize, get acc and report
                if loss is not None:
                    if args.gradient_accumulation_steps > 1:
                        loss = loss / args.gradient_accumulation_steps
                    if args.fp16:
                        with amp.scale_loss(loss, optimizer) as scaled_loss:
                            scaled_loss.backward()
                    else:
                        loss.backward()

                    total_loss += loss.mean().item()

                    if (step_idx + 1) % args.gradient_accumulation_steps == 0:
                        if args.fp16:
                            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                        else:
                            torch.nn.utils.clip_grad_norm_(target_encoder.parameters(), args.max_grad_norm)

                        optimizer.step()
                        scheduler.step()  # Update learning rate schedule
                        target_encoder.zero_grad()

                    pbar.set_description(
                        f"Ep {ep_idx+1} Tr loss: {loss.mean().item():.2f}, acc: {sum(accs)/len(accs):.3f}"
                    )

                #import pdb; pdb.set_trace()
                if accs is not None:
                    total_accs += accs
                    total_accs_k += [len(tgt) > 0 for tgt in p_tgts_t]
                else:
                    total_accs += [0.0]*len(p_tgts_t)
                    total_accs_k += [0.0]*len(p_tgts_t)
                batch = tuple(t.to("cpu") for t in batch)

        step_idx += 1
        logger.info(
            f"Avg train loss ({step_idx} iterations): {total_loss/step_idx:.2f} | train " +
            f"acc@1: {sum(total_accs)/len(total_accs):.3f} | acc@{args.top_k}: {sum(total_accs_k)/len(total_accs_k):.3f}"
        )

        # Evaluation
        new_args = copy.deepcopy(args)
        new_args.top_k = 10
        new_args.save_pred = False
        new_args.test_path = args.dev_path
        dev_em, dev_f1, dev_emk, dev_f1k = evaluate_type(new_args, mips, target_encoder, tokenizer)
        logger.info(f"Develoment set acc@1: {dev_em:.3f}, f1@1: {dev_f1:.3f}")

        # Save best model
        if dev_em > best_acc:
            best_acc = dev_em
            save_path = args.output_dir
            if not os.path.exists(save_path):
                os.makedirs(save_path)
            target_encoder.save_pretrained(save_path)
            logger.info(f"Saved best model with acc {best_acc:.3f} into {save_path}")
            global_acc[-1] = [ep_idx, sum(total_accs)/len(total_accs), best_acc]

        if (ep_idx + 1) % 1 == 0:
            logger.info('Updating pretrained encoder')
            # pretrained_encoder = copy.deepcopy(target_encoder)

    print()
    logger.info(f"Best model has acc {best_acc:.3f} saved as {save_path}")

def get_top_phrases(mips, q_ids, questions, answers, titles, query_encoder, tokenizer, batch_size, args):
    # Search
    step = batch_size
    phrase_idxs = []
    search_fn = mips.search
    query2vec = get_query2vec(
        query_encoder=query_encoder, tokenizer=tokenizer, args=args, batch_size=batch_size
    )
    for q_idx in tqdm(range(0, len(questions), step)):
        outs = query2vec(questions[q_idx:q_idx+step])
        start = np.concatenate([out[0] for out in outs], 0)
        end = np.concatenate([out[1] for out in outs], 0)
        query_vec = np.concatenate([start, end], 1)

        outs = search_fn(
            query_vec,
            q_texts=questions[q_idx:q_idx+step], nprobe=args.nprobe,
            top_k=args.top_k, return_idxs=True,
            max_answer_length=args.max_answer_length, aggregate=args.aggregate, agg_strat=args.agg_strat,
        )
        yield (
            q_ids[q_idx:q_idx+step], questions[q_idx:q_idx+step], answers[q_idx:q_idx+step],
            titles[q_idx:q_idx+step], outs
        )

def annotate_phrase_vecs(mips, q_ids, questions, answers, titles, phrase_groups, args):
    assert mips is not None
    batch_size = len(answers)

    # Phrase groups are in size of [batch, top_k, values]
    # phrase_groups = [[(
    #     out_['doc_idx'], out_['start_idx'], out_['end_idx'], out_['answer'],
    #     out_['start_vec'], out_['end_vec'], out_['context'], out_['title'])
    #     for out_ in out] for out in outs
    # ]
    dummy_group = {
        'doc_idx': -1,
        'start_idx': 0, 'end_idx': 0,
        'answer': '',
        'start_vec': np.zeros(768),
        'end_vec': np.zeros(768),
        'context': '', 'title': ['']
    }

    # Pad phrase groups (two separate top-k coming from start/end, so pad with top_k*2)
    for b_idx, phrase_idx in enumerate(phrase_groups):
        while len(phrase_groups[b_idx]) < args.top_k*2:
            phrase_groups[b_idx].append(dummy_group)
        assert len(phrase_groups[b_idx]) == args.top_k*2

    # Flatten phrase groups
    flat_phrase_groups = [phrase for phrase_group in phrase_groups for phrase in phrase_group]
    doc_idxs = [int(phrase_group['doc_idx']) for phrase_group in flat_phrase_groups]
    start_vecs = [phrase_group['start_vec'] for phrase_group in flat_phrase_groups]
    end_vecs = [phrase_group['end_vec'] for phrase_group in flat_phrase_groups]

    # stack vectors
    start_vecs = np.stack(start_vecs)
    end_vecs = np.stack(end_vecs)
    zero_mask = np.array([[1] if doc_idx >= 0 else [0] for doc_idx in doc_idxs])
    start_vecs = start_vecs * zero_mask
    end_vecs = end_vecs * zero_mask

    # Reshape
    start_vecs = np.reshape(start_vecs, (batch_size, args.top_k*2, -1))
    end_vecs = np.reshape(end_vecs, (batch_size, args.top_k*2, -1))

    # Dummy targets
    targets = [[None for phrase in phrase_group] for phrase_group in phrase_groups]
    p_targets = [[None for phrase in phrase_group] for phrase_group in phrase_groups]

    # TODO: implement dynamic label_strategy based on the task name (label_strat = dynamic)
    # TODO: implement both doc and paragraph as strageties

    # Annotate for L_phrase
    if 'phrase' in args.label_strat.split(','):
        match_fns = [
            drqa_regex_match_score if args.regex or ('trec' in q_id.lower()) else drqa_exact_match_score for q_id in q_ids
        ]
        targets = [
            [drqa_metric_max_over_ground_truths(match_fn, phrase['answer'], answer_set) for phrase in phrase_group]
            for phrase_group, answer_set, match_fn in zip(phrase_groups, answers, match_fns)
        ]
        targets = [[ii if val else None for ii, val in enumerate(target)] for target in targets]

    # Annotate for L_doc
    if 'doc' in args.label_strat.split(','):
        p_targets = [
            [any(phrase['title'][0].lower() == tit.lower() for tit in title) for phrase in phrase_group]
            for phrase_group, title in zip(phrase_groups, titles)
        ]
        p_targets = [[ii if val else None for ii, val in enumerate(target)] for target in p_targets]
    elif "paragraph" in args.label_strat.split(","):
        p_t = [[any([answer == phrase["context"] for answer in answers[i]]) for phrase in phrase_group] for i, phrase_group in enumerate(phrase_groups)]
        p_t = [[ii if val else None for ii, val in enumerate(target)] for target in p_t]
        p_targets = p_t

    return start_vecs, end_vecs, targets, p_targets

if __name__ == '__main__':
    # See options in densephrases.options
    options = Options()
    options.add_model_options()
    options.add_index_options()
    options.add_retrieval_options()
    options.add_data_options()
    options.add_qsft_options()
    options.add_question_type_options()
    args = options.parse()

    # Seed for reproducibility
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    def nice_print(results, print_fn):
        for re in results:
            print_fn(", ".join(map(str,re)))

    if args.run_mode == 'train_query':
        # Train

        device = 'cuda' if args.cuda else 'cpu'
        pretrained_encoder, tokenizer, _ = load_encoder(device, args)
        mips = load_phrase_index(args)

        # --- used in regularization ---
        # output_dir = args.output_dir
        # dropout_sweep = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5]
        # weight_decay_sweep = [0.01, 0.03, 0.05, 0.07, 0.1, 0.2]
        # lr_sweep = [1e-5, 3e-5, 5e-5, 7e-5, 1e-4, 2e-4]
        # layer_freeze_sweep = [1,2,3,4,5,6,7,8,9,10,11]
        # layer_freeze_sweep = [10,11]

        # for i in layer_freeze_sweep:
        #     logger.info("\n------------------\n")
        #     logger.info(f"Starting training for layer freeze = {i}")

        #     new_args = copy.deepcopy(args)

        #     # edit hyperparameter here
        #     new_args.num_layers_freeze = i
        #     #new_args.weight_decay = i
        #     #new_args.learning_rate = i

        #     # for changing dropout
        #     #new_args.config_name = f"/n/fs/nlp-hyen/DensePhrases/custom_config/dropout-0.1"

        #     # edit output dir
        #     new_args.output_dir = output_dir + f"-{i}-freeze"

        #     train_query_encoder(new_args, mips)
        train_query_encoder(args, mips)
        #     # Eval
        #     new_args.query_encoder_path = new_args.output_dir
        #     logger.info(f"Evaluating {new_args.query_encoder_path}")
        #     new_args.config_name = new_args.output_dir
        #     new_args.top_k = 10
        #     eval_em = evaluate(new_args, mips)
        #     global_acc[-1].append(eval_em)
        # logger.info(f"Global Acc (Train, Dev, Eval): {global_acc}")
        # nice_print(global_acc, print)

        # Eval
        args.load_dir = args.output_dir
        logger.info(f"Evaluating {args.load_dir}")
        args.top_k = 10
        evaluate_type(args, mips)

    else:
        raise NotImplementedError


"""
miscelaneous functions

def human_readable_bytes(byt):
def path_up_to_wildcard(full_path):
def parse_url(url, config, use_local=False, has_wildcards=False):
def user_from_config(config, host):
def get_time_string(seconds):
"""

# numpy is faster, but we don't use enough to warrant the extra dependencies
#from numpy import log, power, abs

# use built in python functions
import logging
import getpass
import os
import re
from math import log
from collections import namedtuple
from datetime import datetime, timezone
power = pow
LOGGER = logging.getLogger(name='util')

LN_BASE = log(power(1024, 1/3))
def human_readable_bytes(byt):
    """ fixed version of https://stackoverflow.com/a/17754143/663466
     hybrid of https://stackoverflow.com/a/10171475/2595465
      with https://stackoverflow.com/a/5414105/2595465  """
    # return bytes if small
    if byt <= 99:
        return str(int(byt))
    magnitude = int(log(abs(byt)) / LN_BASE)
    if magnitude > 19:
        float_fmt = '%i'
        illion = 20 // 3
    else:
        mag3 = (magnitude+1) % 3
        float_fmt = '%' + str(mag3) + "." + str(3-mag3) + 'f'
        illion = (magnitude + 1) // 3
    format_str = float_fmt + ['', 'K', 'M', 'G', 'T', 'P', 'E'][illion]
    return (format_str % (byt * 1.0 / (1024 ** illion))).lstrip('0')

def path_up_to_wildcard(full_path):
    """ If a given path has a wildcard placeholder ( eg {sample} ),
    return the last directory before that point """
    path_fragment = full_path.split('{')[0]
    if path_fragment == full_path:
        return full_path
    if path_fragment.endswith(os.path.pathsep):
        return path_fragment[:-1]
    return os.path.dirname(path_fragment)

URL_REXP = re.compile(r'^([A-Za-z]+)://(?:([^/@]+)@)?([^/]*)(/.+)$')
Remote = namedtuple('Remote', ['protocol', 'user', 'host', 'path'])

def parse_url(url, config, use_local=False, has_wildcards=False):
    """ check if the string is a url or simple path
    return None if it's a path
    return named tuple with (protocol, user, host, path) if its a URL
    """

    ## Check 1: is it a full formed URL? EG:
    #   SFTP://server.com/path/to/file
    #   file:///local/path
    #   SCP://user@host.dom/some/path
    match = URL_REXP.search(url)
    if match:
        remote = Remote(*match.groups())
        if remote.user is None:
            user = user_from_config(config, remote.host)
            remote = Remote(remote.protocol, user,
                            remote.host, remote.path)
        if remote.protocol.lower == 'file':
            if len(remote.host) > 0:
                raise Exception("file URL should have no host name")
            return None
        return remote

    # skip check 2 if file exists and we're OK using local files
    if use_local:
        if os.path.exists(path_up_to_wildcard(url) \
                          if has_wildcards else url):
            return None

    ## check 2: user configured remote maps
    for custom_patterns in config \
                           .get('remote', {}) \
                           .get('mappings', []):
        try:
            mnt_rexp = re.compile(custom_patterns['pattern'])
            host_repl = custom_patterns['host_repl']
            path_repl = custom_patterns['path_repl']
        except KeyError:
            LOGGER.error("custom patterns must contain: pattern, host_repl, "
                         "and path_repl")
            raise
        except:
            LOGGER.error("re cannot compile custom pattern: " +
                         custom_patterns['pattern'])
            raise

        LOGGER.debug("Checking remote pattern: %r", custom_patterns['pattern'])

        if not mnt_rexp.search(url):
            # skip to next pattern if this doesn't match
            continue

        try:
            source_path = mnt_rexp.sub(path_repl, url)
        except:
            LOGGER.error("re cannot understand replacement expression " +
                         path_repl)
            raise
        try:
            host = mnt_rexp.sub(host_repl, url)
        except:
            LOGGER.error("re cannot understand replacement expression " +
                         host_repl)
            raise
        user = user_from_config(config, host)

        LOGGER.debug("INFERRED URL SFTP://%s@%s%s", user, host, source_path)
        return Remote('SFTP', user, host, source_path)

    ## 3: just a regular, local file
    # we ge here if there was no match above
    return None

def user_from_config(config, host):
    """ get username from config for this host. Fall back to local username """
    local_user = getpass.getuser()
    LOGGER.debug("config['remote']: %r", config.get('remote',{}))
    default_user = config.get('remote', {})\
                         .get('SFTP', {}) \
                         .get('default', {}) \
                         .get('username', local_user)
    user = config.get('remote', {}).get('SFTP', {}).get(host, {}).get('username', default_user)
    return user

def get_time_string(seconds):
    """ return a formatted time string """
    return datetime.fromtimestamp(seconds,
                                  timezone.utc) \
                    .astimezone() \
                    .strftime("%Y-%m-%d %H:%M:%S (%Z)")


# coding=utf-8
from __future__ import print_function
import multiprocessing
import math
import itertools
import re
import sys, os
import six
import torch
import torch.nn as nn
from collections import OrderedDict

import numpy as np

from components.evaluator import CachedExactMatchEvaluator
from common.registerable import Registrable
from common.savable import Savable
from datasets.conala.conala_eval import tokenize_for_bleu_eval
from datasets.conala import evaluator as conala_evaluator
from model import utils
from components.dataset import Example

if six.PY3:
    from asdl import Python3TransitionSystem
else:
    from asdl import PythonTransitionSystem

import xgboost as xgb

# shared across processes for multi-processed reranking
_examples = None
_decode_results = None
_evaluator = None
_ranker = None

def _rank_worker(param):
    score = _ranker.compute_rerank_performance(_examples, _decode_results, fast_mode=True, evaluator=_evaluator, param=param)
    return param, score

def _rank_segment_worker(param_space):
    best_score = 0.
    best_param = None
    print('[Child] New parameter segments [%s ~ %s] (%d entries)' % (param_space[0], param_space[-1], len(param_space)), file=sys.stderr)
    for param in param_space:
        score = _ranker.compute_rerank_performance(_examples, _decode_results, fast_mode=True, evaluator=_evaluator, param=np.array(param))
        if score > best_score:
            print('[Child] New param=%s, score=%.4f' % (param, score), file=sys.stderr)
            best_param = param
            best_score = score

    return best_param, best_score

class RerankingFeature(object):
    @property
    def feature_name(self):
        raise NotImplementedError

    @property
    def is_batched(self):
        raise NotImplementedError

    def get_feat_value(self, example, hyp, **kwargs):
        raise NotImplementedError

@Registrable.register('parser_score')
class ParserScore(RerankingFeature):
    @property
    def feature_name(self):
        return 'parser_score'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        return float(hyp.score) * len(hyp.actions)

@Registrable.register('normalized_parser_score_by_action')
class NormalizedParserScoreByAction(RerankingFeature):
    def __init__(self):
        pass

    @property
    def feature_name(self):
        return 'normalized_parser_score'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        # return float(hyp.score) / len(kwargs['transition_system'].tokenize_code(hyp.code))
        return float(hyp.score)

@Registrable.register('normalized_parser_score')
class NormalizedParserScore(RerankingFeature):
    def __init__(self):
        pass

    @property
    def feature_name(self):
        return 'normalized_parser_score'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        return float(hyp.score) * len(hyp.actions) / HypCodeTokensCount().get_feat_value(example, hyp, **kwargs)

@Registrable.register('word_cnt')
class HypCodeTokensCount(RerankingFeature):
    @property
    def feature_name(self):
        return 'word_cnt'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        # TODO: this is dataset specific, we should fix it
        if hasattr(hyp, 'decanonical_code_tokens'):
            # we use the tokenization for BLEU calculation
            code_tokens = ['#NEWLINE#' if c == '\n' else c for c in hyp.decanonical_code_tokens]
            return float(len(code_tokens))
        # else:
        #     return len(hyp.actions)
        elif six.PY2 and isinstance(kwargs['transition_system'], PythonTransitionSystem):
            code_tokens = [c.replace('\r\n', '#NEWLINE#').replace('\r', '#NEWLINE#').replace('\n', '#NEWLINE#')
                           for c in kwargs['transition_system'].tokenize_code(hyp.code)]
            # remove consecutive spaces
            code_tokens = re.sub(r'\s+', ' ', ' '.join(code_tokens)).strip().split(' ')
            code_tokens = list(filter(lambda x: len(x) > 0, code_tokens))
            return float(len(code_tokens))

        return len(kwargs['transition_system'].tokenize_code(hyp.code))

@Registrable.register('is_2nd_hyp_and_margin_with_top_hyp')
class IsSecondHypAndScoreMargin(RerankingFeature):
    def __init__(self):
        pass

    @property
    def feature_name(self):
        return 'is_2nd_hyp_and_margin_with_top_hyp'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        if kwargs['hyp_id'] == 1:
            return kwargs['all_hyps'][0].score - hyp.score
        return 0.

@Registrable.register('is_2nd_hyp_and_paraphrase_score_margin_with_top_hyp')
class IsSecondHypAndParaphraseScoreMargin(RerankingFeature):
    def __init__(self):
        pass

    @property
    def feature_name(self):
        return 'is_2nd_hyp_and_paraphrase_score_margin_with_top_hyp'

    @property
    def is_batched(self):
        return False

    def get_feat_value(self, example, hyp, **kwargs):
        if kwargs['hyp_id'] == 1:
            return hyp.rerank_feature_values['paraphrase_score'] - kwargs['all_hyps'][0].rerank_feature_values['paraphrase_score']
        return 0.

@Registrable.register('reranker')
class Reranker(Savable):
    def __init__(self, features, parameter=None, transition_system=None):
        self.features = []
        self.transition_system = transition_system
        self.feat_map = OrderedDict()
        self.batched_features = OrderedDict()

        for feat in features:
            self._add_feature(feat)

        if parameter is not None:
            self.parameter = parameter
        else:
            self.parameter = np.zeros(self.feature_num)

    def _add_feature(self, feature):
        self.features.append(feature)
        self.feat_map[feature.feature_name] = feature

        if feature.is_batched:
            self.batched_features[feature.feature_name] = feature

    def get_initial_reranking_feature_values(self, example, hyp, **kwargs):
        """Given a hypothesis, compute its reranking feature"""
        feat_values = OrderedDict()
        for feat_name, feat in self.feat_map.items():
            if not feat.is_batched:
                feat_val = feat.get_feat_value(example, hyp, **kwargs)
            else:
                feat_val = float('inf')

            feat_values[feat_name] = feat_val

        return feat_values

    def rerank_hypotheses(self, examples, decode_results, fast_mode=False):
        """rerank the hypotheses using the current model parameter"""
        self.filter_hyps_and_initialize_features(examples, decode_results)

        param = self.parameter

        sorted_decode_results = []
        for example, hyps in zip(examples, decode_results):
            if hyps:
                new_hyp_scores = [self.get_rerank_score(hyp, param=param) for hyp in hyps]
                for score, hyp in zip(new_hyp_scores, hyps):
                    hyp.rerank_score = score

                if fast_mode:
                    best_hyp_idx = np.argmax(new_hyp_scores)
                    best_hyp = hyps[best_hyp_idx]
                    sorted_decode_results.append([best_hyp])
                else:
                    sorted_decode_results.append([hyps[i] for i in np.argsort(new_hyp_scores)[::-1]])
            else:
                sorted_decode_results.append([])

        return sorted_decode_results

    def initialize_rerank_features(self, examples, decode_results):
        hyp_examples = []
        # print('initializing features...', file=sys.stderr)
        for example, hyps in zip(examples, decode_results):
            for hyp_id, hyp in enumerate(hyps):
                hyp_example = Example(idx=None,
                                      src_sent=example.src_sent,
                                      tgt_code=hyp.code,
                                      tgt_actions=None,
                                      tgt_ast=None)
                hyp_examples.append(hyp_example)
                # hyp.tokenized_code = len(self.transition_system.tokenize_code(hyp.code))
                # hyp.code_token_count = len(hyp.code.split(' '))

                feat_vals = OrderedDict()
                hyp.rerank_feature_values = feat_vals

        for batch_examples in utils.batch_iter(hyp_examples, batch_size=128):
            for feat_name, feat in self.batched_features.items():
                batch_example_scores = feat.score(batch_examples).data.cpu().tolist()
                for i, e in enumerate(batch_examples):
                    setattr(e, feat_name, batch_example_scores[i])

        e_ptr = 0
        for example, hyps in zip(examples, decode_results):
            for hyp in hyps:
                for feat_name, feat in self.batched_features.items():
                    hyp.rerank_feature_values[feat_name] = getattr(hyp_examples[e_ptr], feat_name)
                e_ptr += 1

        for example, hyps in zip(examples, decode_results):
            for hyp_id, hyp in enumerate(hyps):
                for feat_name, feat in self.feat_map.items():
                    if not feat.is_batched:
                        feat_val = feat.get_feat_value(example, hyp,
                                                       hyp_id=hyp_id, all_hyps=hyps,
                                                       transition_system=self.transition_system)
                        hyp.rerank_feature_values[feat_name] = feat_val

    def get_rerank_score(self, hyp, param):
        raise NotImplementedError

    def _filter_hyps(self, decode_results, is_valid_hyp):
        for i in range(len(decode_results)):
            valid_hyps = []
            for hyp in decode_results[i]:
                if is_valid_hyp(hyp):
                    valid_hyps.append(hyp)

            decode_results[i] = valid_hyps

    def filter_hyps_and_initialize_features(self, examples, decode_results):
        # print(decode_results[0][0].code)
        if not hasattr(decode_results[0][0], 'rerank_feature_values'):
            # print('initializing rerank features for hypotheses...', file=sys.stderr)

            def is_valid_hyp(hyp):
                try:
                    self.transition_system.tokenize_code(hyp.code)
                    if hyp.code:
                        return True
                except:
                    return False

                return False

            self._filter_hyps(decode_results, is_valid_hyp)

            self.initialize_rerank_features(examples, decode_results)

    def compute_rerank_performance(self, examples, decode_results, evaluator=CachedExactMatchEvaluator(),
                                   param=None, fast_mode=False, verbose=False, args=None):
        self.filter_hyps_and_initialize_features(examples, decode_results)

        if param is None:
            param = self.parameter

        sorted_decode_results = []
        for example, hyps in zip(examples, decode_results):
            if hyps:
                new_hyp_scores = [self.get_rerank_score(hyp, param=param) for hyp in hyps]
                best_hyp_idx = np.argmax(new_hyp_scores)
                best_hyp = hyps[best_hyp_idx]

                if fast_mode:
                    sorted_decode_results.append([best_hyp])
                else:
                    sorted_decode_results.append([hyps[i] for i in np.argsort(new_hyp_scores)[::-1]])
            else:
                sorted_decode_results.append([])

            if verbose:
                gold_standard_idx = [i for i, hyp in enumerate(hyps) if hyp.is_correct]
                if gold_standard_idx and gold_standard_idx[0] != best_hyp_idx:
                    gold_standard_idx = gold_standard_idx[0]
                    print('Utterance: %s' % ' '.join(example.src_sent), file=sys.stderr)
                    print('Gold hyp id: %d' % gold_standard_idx, file=sys.stderr)
                    for _i, hyp in enumerate(hyps):
                        print('Hyp %d: %s ||| score: %f ||| final score: %f' % (_i,
                                                                                hyp.code,
                                                                                hyp.score,
                                                                                self.get_rerank_score(hyp, param=param)),
                              file=sys.stderr)
                        print('\t%s' % hyp.rerank_feature_values, file=sys.stderr)

        metric = evaluator.evaluate_dataset(examples, sorted_decode_results, fast_mode=fast_mode, args=args)

        return metric

    def train(self, examples, decode_results, initial_performance=0., metric='accuracy'):
        raise NotImplementedError

    @property
    def feature_num(self):
        return len(self.features)

    def __getattr__(self, item):
        if item in self.feat_map:
            return self.feat_map.get(item)
        raise ValueError

    def save(self, path):
        dir_name = os.path.dirname(path)
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)

        feature_names = []
        for feat in self.features:
            if isinstance(feat, nn.Module):
                feat.save(os.path.join(path + '.%s' % feat.feature_name))
            feature_names.append(feat.feature_name)

        params = {
            'parameter': self.parameter,
            'feature_names': feature_names,
            'transition_system': self.transition_system
        }

        torch.save(params, path)

    @classmethod
    def load(cls, model_path, cuda=False):
        print("loading reranker...")
        params = torch.load(model_path, map_location=lambda storage, loc: storage)
        feature_names = params['feature_names']
        features = []
        for feat_name in feature_names:
            feat_cls = Registrable.registered_components[feat_name]
            if issubclass(feat_cls, Savable):
                feat_inst = feat_cls.load(model_path + '.%s' % feat_name, cuda=cuda)
                feat_inst.eval()
            else:
                feat_inst = feat_cls()
            features.append(feat_inst)

        reranker = cls(features, params['parameter'], params['transition_system'])

        return reranker

class LinearReranker(Reranker):
    """a reranker using linear features"""

    def get_rerank_score(self, hyp, param=None):
        if param is None:
            param = self.parameter

        score = sum(feat_weight * hyp.rerank_feature_values[feat_name] for feat_name, feat_weight in param.items())

        return score

    def generate_nbest_list(self, examples, decode_results, nbest_file_name, target_tokenizer, hyp_tokenizer):
        f_src = open(nbest_file_name + '.src', 'w')
        f_tgt = open(nbest_file_name + '.tgt', 'w')
        f_hyp = open(nbest_file_name + '.hyp', 'w')

        self.filter_hyps_and_initialize_features(examples, decode_results)

        for e_id, (example, hyp_list) in enumerate(zip(examples, decode_results)):
            f_src.write(' '.join(example.src_sent) + '\n')
            f_tgt.write(' '.join(target_tokenizer(example)) + '\n')

            if not hyp_list:
                hyp_str = "{e_id} ||| pass ||| 0. |||".format(e_id=e_id)
                for feat_name in self.feat_map:
                    hyp_str += ' {}=0.0'.format(feat_name)
                hyp_str += '\n'
                f_hyp.write(hyp_str)

                continue

            # new_hyp_scores = [self.get_rerank_score(hyp) for hyp in hyp_list]

            # hyp_ranks = np.argsort(new_hyp_scores)[::-1]

            # for i, hyp_id in enumerate(hyp_ranks):
            for hyp_id in range(len(hyp_list)):
                hyp = hyp_list[hyp_id]
                code_tokens = hyp_tokenizer(hyp)
                # reranker_score = new_hyp_scores[hyp_id]

                # code_token_count={code_token_count_feat.get_feat_value(example, hyp)}
                hyp_str = "{e_id} ||| {code_tokens} ||| {hyp_score} |||".format(e_id=e_id, code_tokens=' '.join(code_tokens), hyp_score=hyp.score)
                for feat_name in self.parameter.keys():
                    hyp_str += ' {feat_name}={feat_val}'.format(feat_name=feat_name, feat_val=hyp.rerank_feature_values[feat_name])
                hyp_str += '\n'
                f_hyp.write(hyp_str)

        f_src.close()
        f_tgt.close()
        f_hyp.close()

    def train(self, examples, decode_results, initial_performance=0., metric='accuracy'):
        raise NotImplementedError

class GridSearchReranker(Reranker):
    """Grid search reranker"""

    def get_rerank_score(self, hyp, param):
        feat_vals = np.array(list(hyp.rerank_feature_values.values()))
        score = hyp.score + np.dot(param, feat_vals)
        # score = np.dot(param, feat_vals)

        return score

    def train(self, examples, decode_results, evaluator=CachedExactMatchEvaluator(), initial_performance=0.):
        """optimize the ranker on a dataset using grid search"""
        best_score = initial_performance
        best_param = np.zeros(self.feature_num)

        param_space = (np.array(p) for p in itertools.combinations(np.arange(0, 3.01, 0.01), self.feature_num))

        for param in param_space:
            score = self.compute_rerank_performance(examples, decode_results, fast_mode=True, evaluator=evaluator, param=param)
            if score > best_score:
                print('New param=%s, score=%.4f' % (param, score), file=sys.stderr)
                best_param = param
                best_score = score

        self.parameter = best_param

    def train_multiprocess(self, examples, decode_results, evaluator=CachedExactMatchEvaluator(), initial_performance=0., num_workers=8):
        """optimize the ranker on a dataset using grid search"""
        best_score = initial_performance
        best_param = np.zeros(self.feature_num)

        self.initialize_rerank_features(examples, decode_results)

        print('generating parameter list', file=sys.stderr)
        param_space = [p for p in itertools.combinations(np.arange(0, 2.03, 0.02), self.feature_num)]
        print('generating parameter list done', file=sys.stderr)

        global _examples
        _examples = examples
        global _decode_results
        _decode_results = decode_results
        global _evaluator
        _evaluator = evaluator
        global _ranker
        _ranker = self

        def _norm(_param):
            return sum(p ** 2 for p in _param)

        with multiprocessing.Pool(processes=num_workers) as pool:
            # segment the parameter space
            segment_size = int(len(param_space) / num_workers / 5)
            param_space_segments = []
            ptr = 0
            while ptr < len(param_space):
                param_space_segments.append(param_space[ptr: ptr + segment_size])
                ptr += segment_size
            print('generated %d parameter segments' % len(param_space_segments), file=sys.stderr)

            results = pool.imap_unordered(_rank_segment_worker, param_space_segments)

            for param, score in results:
                if score > best_score or score == best_score and _norm(param) < _norm(best_param):
                    print('[Main] New param=%s, score=%.4f' % (param, score), file=sys.stderr)
                    best_param = param
                    best_score = score

        self.parameter = best_param

class XGBoostReranker(Reranker):
    def __init__(self, features, transition_system=None):
        super(XGBoostReranker, self).__init__(features, transition_system=transition_system)

        params = {'objective': 'rank:ndcg', 'learning_rate': .1,
                  'gamma': 5.0, 'min_child_weight': 0.1,
                  'max_depth': 4, 'n_estimators': 5}

        self.ranker = xgb.sklearn.XGBRanker(**params)

    def get_feature_matrix(self, decode_results, train=False):
        x, y, group = [], [], []

        for hyps in decode_results:
            if hyps:
                for hyp in hyps:
                    label = 1 if hyp.is_correct else 0
                    feat_vec = np.array([hyp.score] + [v for v in hyp.rerank_feature_values.values()])
                    x.append(feat_vec)
                    y.append(label)
                group.append(len(hyps))

        x = np.stack(x)
        y = np.array(y)

        # if train:
        #     self.scaler = preprocessing.StandardScaler().fit(x)
        #     x = self.scaler.transform(x)
        # else:
        #     x = self.scaler.transform(x)

        return x, y, group

    def get_rerank_score(self, hyp, param):
        x, y, group = self.get_feature_matrix([[hyp]])
        y = self.ranker.predict(x)

        return y[0]

    def train(self, examples, decode_results, evaluator=CachedExactMatchEvaluator(), initial_performance=0.):
        self.initialize_rerank_features(examples, decode_results)

        train_x, train_y, group_train = self.get_feature_matrix(decode_results, train=True)
        self.ranker.fit(train_x, train_y, group_train)

        train_acc = self.compute_rerank_performance(examples, decode_results, fast_mode=True, evaluator=evaluator)
        print('Dev acc: %f' % train_acc, file=sys.stderr)