# -*- coding: utf-8 -*-
"""
Custom loss functions for the autoencoder on the tensorflow backend.
"""
import tensorflow as tf
from keras.backend.common import epsilon
from keras.backend.tensorflow_backend import _to_tensor
from keras import backend as K
 
#return a dict of the custom loss functions to pass to load_model
def get_custom_objects():
    custom_objects = {}
    
    custom_objects['mean_squared_error_poisson'] = mean_squared_error_poisson
    custom_objects['msep_squared'] = msep_squared
    custom_objects['msep_log'] = msep_log
    custom_objects["cat_cross_inv"]=cat_cross_inv
    custom_objects["wasserstein_loss"]=wasserstein_loss
    
    return custom_objects

#like mse, except every bin is weighted with the probability 
#of it not being generated by poisson
def mean_squared_error_poisson(y_true, y_pred):
    #the mean of y_true is the approximated poisson expectation value
    expec = tf.multiply(tf.ones_like(y_true), tf.reduce_mean(y_true))
    #The probability to get a certain bin of y_true from poisson
    #ranges from   0: definitly not Poisson    to 1: definitly poisson
    poisson_factor = tf.exp(-1*expec) * tf.pow(expec, y_true) / tf.exp(tf.lgamma(y_true+1))
    #return weighted mean squared error
    return tf.reduce_mean(tf.squared_difference(y_true, y_pred) * (1-poisson_factor))


#like mse, except every bin is weighted with the square probability 
#of it not being generated by poisson
def msep_squared(y_true, y_pred):
    #the mean of y_true is the approximated poisson expectation value
    expec = tf.multiply(tf.ones_like(y_true), tf.reduce_mean(y_true))
    #The probability to get a certain bin of y_true from poisson
    #ranges from   0: definitly not Poisson    to 1: definitly poisson
    poisson_factor = tf.exp(-1*expec) * tf.pow(expec, y_true) / tf.exp(tf.lgamma(y_true+1))
    #return weighted mean squared error
    return tf.reduce_mean(tf.squared_difference(y_true, y_pred) * tf.square(1-poisson_factor))

#like mse, except every bin is weighted with -1 times the log of the probability 
#of it being generated by poisson
def msep_log(y_true, y_pred):
    #the mean of y_true is the approximated poisson expectation value
    expec = tf.multiply(tf.ones_like(y_true), tf.reduce_mean(y_true))
    #The probability to get a certain bin of y_true from poisson
    #ranges from   0: definitly not Poisson    to 1: definitly poisson
    poisson_factor = tf.exp(-1*expec) * tf.pow(expec, y_true) / tf.exp(tf.lgamma(y_true+1))
    #return weighted mean squared error
    return tf.clip_by_value(tf.reduce_mean(tf.squared_difference(y_true, y_pred) *(-1/100)*tf.log(poisson_factor)),0,100)


def cat_cross_inv(y_true, y_pred):
    #like the normal categorical crossentropy, but labels are reversed!
    #So this is how wrong the network was
    y_true=1-y_true
    axis=-1
    y_pred /= tf.reduce_sum(y_pred, axis, True)
    # manual computation of crossentropy
    _epsilon = _to_tensor(epsilon(), y_pred.dtype.base_dtype)
    y_pred = tf.clip_by_value(y_pred, _epsilon, 1. - _epsilon)
    return - tf.reduce_sum(y_true * tf.log(y_pred), axis)


def wasserstein_loss(y_true, y_pred):
    """Calculates the Wasserstein loss for a sample batch.
    The Wasserstein loss function is very simple to calculate. In a standard GAN, the discriminator
    has a sigmoid output, representing the probability that samples are real or generated. In Wasserstein
    GANs, however, the output is linear with no activation function! Instead of being constrained to [0, 1],
    the discriminator wants to make the distance between its output for real and generated samples as large as possible.
    The most natural way to achieve this is to label generated samples -1 and real samples 1, instead of the
    0 and 1 used in normal GANs, so that multiplying the outputs by the labels will give you the loss immediately.
    Note that the nature of this loss means that it can be (and frequently will be) less than 0."""
    return K.mean(y_true * y_pred, axis=(1,2))


if __name__=="__main__":
    import numpy as np
    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        #a=np.random.randint(0,4,10)
        a=np.random.rand(5,2,1)
        b=np.random.rand(5,2,1)
        x = tf.constant(a, shape=a.shape, dtype="float32")
        y = tf.constant(b, shape=b.shape, dtype="float32")
        print(sess.run(wasserstein_loss(x,y)))

