<gh_stars>0
from .network import network_batchnorm, network_standard
import tensorflow as tf
from .config import Config


class Classifier(object):

    def __init__(self, num_label):
        config = Config()
        self.params = config.params
        self.model_dir = config.model_dir
        self.params['n_classes'] = num_label


    def get_model_fn(self):
        raise NotImplementedError()


    def loss_fn(self, logits, labels):
        raise NotImplementedError()


    def score_fn(self, logits):
        raise NotImplementedError()


    def pred_fn(self, logits):
        raise NotImplementedError()


    def auc_fn(self, labels, preds):
        auc, update_op = tf.metrics.auc(labels=labels,
                                        predictions=preds,
                                        name='auc')
        return tf.convert_to_tensor(auc), update_op


    def get_estimator(self):

        sess_config = tf.ConfigProto()
        sess_config.gpu_options.allow_growth = True

        estimator = tf.estimator.Estimator(
            model_fn=self.get_model_fn(),
            params=self.params,
            model_dir=self.model_dir,
            config=tf.contrib.learn.RunConfig(session_config=sess_config, save_checkpoints_secs=60))
        return estimator


class CenterlossClassifier(Classifier):

    def __init__(self, num_label):
        config = Config()
        self.params = config.params
        self.model_dir = config.model_dir
        self.batch_size = config.batch_size
        self.params['n_classes'] = num_label


    def get_model_fn(self):
        loss_fn = self.loss_fn
        score_fn = self.score_fn
        pred_fn = self.pred_fn
        auc_fn = self.auc_fn

        def model_fn(features, labels, mode, params):
            if mode == tf.estimator.ModeKeys.TRAIN:
                is_training = True
            else:
                is_training = False

            input_tensor = features['x']
            idx = features['idx']

            if params['batch_norm']:
                logits, bottleneck = network_batchnorm(input_tensor, params, is_training)
            else:
                logits, bottleneck = network_standard(input_tensor, params, is_training)

            centers = tf.get_variable('centers',
                                      [self.params['n_classes'], params['hidden_units'][1]],
                                      dtype=tf.float32,
                                      initializer=tf.random_normal_initializer(),
                                      trainable=False)

            score = score_fn(logits)
            preds = pred_fn(logits)
            predictions = {
                           'bottleneck': bottleneck,
                           'sigmoids': score,
            }

            if mode == tf.estimator.ModeKeys.PREDICT:
                return tf.estimator.EstimatorSpec(mode, predictions=predictions)

            loss, c_update_op = loss_fn(logits, labels, centers, bottleneck)
            sparse_labels = pred_fn(labels)
            auc, auc_update_op = auc_fn(labels, score)

            if mode == tf.estimator.ModeKeys.EVAL:

                metrics = {'auc': (auc, auc_update_op)}

                return tf.estimator.EstimatorSpec(
                    mode, loss=loss, eval_metric_ops=metrics)

            elif mode == tf.estimator.ModeKeys.TRAIN:
                optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
                bn_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                gradients = optimizer.compute_gradients(loss)
                with tf.control_dependencies(bn_update_ops+[c_update_op]):
                    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
                train_ops = tf.group(train_op, auc_update_op, auc)

                logging_hook = tf.train.LoggingTensorHook({'loss':loss,
                                                           'auc':auc,
                                                           'predictions':preds[0:5],
                                                           'labels':sparse_labels[0:5]
                                                            },
                                                          every_n_iter=100)

                return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_ops,
                                                  training_hooks=[logging_hook])

        return model_fn


    def loss_fn(self, logits, labels, centers, bottleneck):
        with tf.variable_scope('loss_definition'):
            with tf.variable_scope('cross_entropy_loss'):
                xentropy_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=labels, logits=logits)
            with tf.variable_scope('center_loss'):
                center_loss, c_update_op = self.get_center_loss(labels, centers, bottleneck, alpha=0.1)
            with tf.variable_scope('regularizer'):
                l2_loss = tf.losses.get_regularization_loss()
            with tf.variable_scope('total_loss'):
                loss = xentropy_loss  + self.params['centerloss_coef'] * center_loss  + self.params['reg_coef'] * l2_loss

        return loss, c_update_op


    def get_center_loss(self, labels, centers, bottleneck):
        raise NotImplementedError()


class MultilabelCenterlossClassifier(CenterlossClassifier):


    def score_fn(self, logits):
        score = tf.sigmoid(logits)
        return score


    def pred_fn(self, logits):
        values, indices = tf.nn.top_k(logits, k=7, sorted=True)
        return indices


    def get_center_loss(self, labels, centers, bottleneck, alpha=0.5):
        """Center loss based on the paper "A Discriminative Feature Learning Approach for Deep Face Recognition"
        (http://ydwen.github.io/papers/WenECCV16.pdf)
        """
        labels = tf.cast(labels, tf.float32)
        one_num = tf.reduce_sum(labels, axis=1)
        one_num = tf.stop_gradient(tf.reshape(one_num, [-1,1]))
        center_batch = self.get_centers_batch(labels, centers, one_num)
        unique_count = tf.reduce_sum(labels, axis=0)
        unique_count = tf.reshape(unique_count, [-1, 1])

        appear_times = tf.matmul(labels, unique_count)
        diff = center_batch - bottleneck
        loss = tf.nn.l2_loss(diff)
        diff = alpha * diff / tf.cast((1 + appear_times), tf.float32)
        diff_matrix = tf.matmul(labels, diff, transpose_a=True)
        update_op = tf.assign_sub(centers, diff_matrix)

        return loss, update_op


    def get_centers_batch(self, labels_float, centers, one_num):
        raise NotImplementedError()


    def get_diff(self, alpha, centers_batch, bottleneck, one_num):
        raise NotImplementedError()


class MultilabelAddCenterlossClassifier(MultilabelCenterlossClassifier):


    def get_centers_batch(self, labels_float, centers, one_num):
        centers_batch = tf.matmul(labels_float, centers)
        return centers_batch

    def get_diff(self, alpha, centers_batch, bottleneck, one_num):
        diff = (1 - alpha) * (centers_batch - bottleneck) / one_num
        return diff
