import tensorflow as tf
import keras as K


def focal_loss(gamma=2., alpha=.5):
    # 3-0.5 precision-low recall-high
    # 3-0.37 precision-high recall-low
    # 3-0.25 precision-high recall-low
    # 2-0.5 precision-low recall-high
    # 2-0.25 precision-high recall-low
    def f_loss(y_true, y_pred):
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
        return - K.backend.sum(alpha * K.backend.pow(1. - pt_1, gamma)
                               * K.backend.log(K.backend.epsilon()+pt_1))\
               - K.backend.sum((1-alpha) * K.backend.pow(pt_0, gamma)
                               * K.backend.log(1. - pt_0 + K.backend.epsilon()))
    return f_loss


def union_loss(gamma=2., alpha=.5):
    def _loss(y_true, y_pred):

        return focal_loss(gamma, alpha)