import tensorflow as tf
from keras import backend as K
from keras.engine.topology import Layer


class CenterLossLayer(Layer):

    def __init__(self, num_classes, feature_dim, alpha_center, **kwargs):
        super().__init__(**kwargs)
        self.alpha_center = alpha_center
        self.num_classes = num_classes
        self.feature_dim = feature_dim

    def build(self, input_shape):
        self.centers = self.add_weight(name='centers',
                                       shape=(self.num_classes, self.feature_dim),
                                       initializer='uniform',
                                       trainable=False)
        # self.counter = self.add_weight(name='counter',
        #                                shape=(1,),
        #                                initializer='zeros',
        #                                trainable=False)  # just for debugging
        super().build(input_shape)

    def call(self, x, mask=None):
        # x[0] is Nx2, x[1] is Nx10 onehot, self.centers is 10x2
        delta_centers = K.dot(K.transpose(x[1]), (K.dot(x[1], self.centers) - x[0]))  # 10x2
        center_counts = K.sum(K.transpose(x[1]), axis=1, keepdims=True) + 1  # 10x1
        delta_centers /= center_counts
        new_centers = self.centers - self.alpha_center * delta_centers
        self.add_update((self.centers, new_centers), x)

        # self.add_update((self.counter, self.counter + 1), x)

        self.result = x[0] - K.dot(x[1], self.centers)
        self.result = K.sum(self.result ** 2, axis=1, keepdims=True)  # / K.dot(x[1], center_counts)
        return self.result  # Nx1

    def compute_output_shape(self, input_shape):
        return K.int_shape(self.result)

    def get_config(self):
        config = {
            'num_classes': self.num_classes,
            'feature_dim': self.feature_dim,
            'alpha_center': self.alpha_center
        }
        base_config = super(CenterLossLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


### custom loss

def zero_loss(y_true, y_pred):
    return 0.5 * K.sum(y_pred, axis=0)


class AMSoftmax(Layer):
    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        self.kernel = None
        super(AMSoftmax, self).__init__(**kwargs)

    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[-1]

        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='glorot_uniform',
                                      trainable=True)
        super(AMSoftmax, self).build(input_shape)

    def call(self, inputs):
        inputs = tf.nn.l2_normalize(inputs, dim=1)  # input_l2norm
        self.kernel = tf.nn.l2_normalize(self.kernel, dim=0)  # W_l2norm

        cosine = K.dot(inputs, self.kernel)  # cos = input_l2norm * W_l2norm
        return cosine

    def compute_output_shape(self, input_shape):
        return input_shape[0], self.output_dim

    def get_config(self):
        config = {
            'output_dim': self.output_dim
        }
        base_config = super(AMSoftmax, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def amsoftmax_loss(y_true, y_pred):
    scale = 30.0
    margin = 0.35

    label = tf.reshape(tf.argmax(y_true, axis=-1), shape=(-1, 1))
    label = tf.cast(label, dtype=tf.int32)  # y
    batch_range = tf.reshape(tf.range(tf.shape(y_pred)[0]), shape=(-1, 1))  # 0~batchsize-1
    indices_of_groundtruth = tf.concat([batch_range, tf.reshape(label, shape=(-1, 1))],
                                       axis=1)  # 2columns vector, 0~batchsize-1 and label
    groundtruth_score = tf.gather_nd(y_pred, indices_of_groundtruth)  # score of groundtruth

    m = tf.constant(margin, name='m')
    s = tf.constant(scale, name='s')

    added_margin = tf.cast(tf.greater(groundtruth_score, m),
                           dtype=tf.float32) * m  # if groundtruth_score>m, groundtruth_score-m
    added_margin = tf.reshape(added_margin, shape=(-1, 1))
    added_embeddingFeature = tf.subtract(y_pred, y_true * added_margin) * s  # s(cos_theta_yi-m), s(cos_theta_j)

    cross_ent = tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=added_embeddingFeature)
    loss = tf.reduce_mean(cross_ent)
    return loss


def confusion_matrix_cell(true_class, pred_class):
    """
    Creates a metric that collects the value for a single cells of the confusion matrix
    :param true_class:
    :param pred_class:
    :return: Keras metric
    """

    def confusion(y_true, y_pred):
        """
        Collects the samples predicted as pred_class where its true class is true_class
        :param y_true:
        :param y_pred:
        :return: the number of predictions as mentioned above
        """
        # Calculate the label from one-hot encoding
        pred_class_label = K.argmax(y_pred, axis=-1)
        true_class_label = K.argmax(y_true, axis=-1)

        # Create a mask representing where the prediction is pred_class and the true class is true_class
        pred_mask = K.equal(pred_class_label, pred_class)
        true_mask = K.equal(true_class_label, true_class)
        mask = tf.logical_and(pred_mask, true_mask)

        # Get the total number of occurences
        occurrences = K.sum(K.cast(mask, 'int32'), axis=0)
        return occurrences

    return confusion


def get_custom_objects():
    return {
        'tf': tf,
        'AMSoftmax': AMSoftmax,
        'amsoftmax_loss': amsoftmax_loss,
        'CenterLossLayer': CenterLossLayer,
        'zero_loss': zero_loss,
        'confusion': confusion_matrix_cell(0, 0)
    }
