import tensorflow as tf


total_epochs = 30
iteration = 14089 // 1
# 128 * 14089 ~ 1,803,460
test_iteration = 10

def center_loss(features, label, alfa, nrof_classes):
    """Center loss based on the paper "A Discriminative Feature Learning Approach for Deep Face Recognition"
       (http://ydwen.github.io/papers/WenECCV16.pdf)
    """
    nrof_features = features.get_shape()[1]
    centers = tf.get_variable('centers', [nrof_classes, nrof_features], dtype=tf.float32,
        initializer=tf.constant_initializer(0), trainable=False)
    label = tf.reshape(label, [-1])
    centers_batch = tf.gather(centers, label)
    diff = (1 - alfa) * (centers_batch - features)
    centers = tf.scatter_sub(centers, label, diff)
    # centers = tf.nn.l2_normalize(centers, 1, 1e-10, name='centers_norm')
    loss = tf.reduce_mean(tf.square(features - centers_batch))
    return loss, centers

def focal_loss(onehot_labels, cls_preds,
                alpha=0.25, gamma=2.0, name=None, scope=None):
    """Compute softmax focal loss between logits and onehot labels
