import os
import sys
import numpy as np
from time import time
import tensorflow as tf
from keras import backend as K

########## weighted binary crossentropy
def create_categorical_crossentropy(weights):
    def weighted_binary_crossentropy(y_true,y_pred):
        b_ce=K.binary_crossentropy(y_true,y_pred)
        weight_vector=y_true*weights[1]+(1-y_true)*weights[0]
        weighted_b_ce=weight_vector*b_ce
        return K.mean(weighted_b_ce)
    return weighted_binary_crossentropy


############ dice loss
def dice_coefficient(y_true,y_pred):
    smoothing_factor=1
    y_true_f=K.flatten(y_true)
    y_pred_f=K.flatten(y_pred)
    intersection=K.sum(y_true_f*y_pred_f)
    return (2.0*intersection+smoothing_factor)/(K.sum(y_true_f)+K.sum(y_pred_f)+smoothing_factor)

########### penalize FN
def balckRegionLoss(y_true,y_pred):
    smoothing_factor=1e-6
    new_y_true = K.ones_like(y_true) - y_true
    y_true_f=K.flatten(new_y_true)
    y_pred_f=K.flatten(y_pred)
    intersection=K.sum(y_true_f*y_pred_f)
    return (2.0*intersection+smoothing_factor)/(K.sum(y_true_f)+K.sum(y_pred_f)+smoothing_factor)

def loss_dice_coefficient_error(y_true,y_pred):
    return 1-dice_coefficient(y_true,y_pred) 


def Tversky_loss(alpha=0.35,beta=0.65):
    def Tversky_loss_inner(y_true,y_pred):
        smooth = 1e-6
        y_true_f=K.flatten(y_true)
        y_pred_f=K.flatten(y_pred)
        intersection=K.sum(y_true_f*y_pred_f)
        tmp = tf.less(y_pred_f,0.5)

        negative_prediction_region = tf.where(tf.less(y_pred_f,0.5),1-y_pred_f,tf.zeros_like(y_pred_f))

        negative_label_region = 1-y_true_f
        FP = K.sum(negative_label_region*y_pred_f)

        FN = K.sum(negative_prediction_region*y_true_f)
        return 1-(intersection+smooth)/(alpha*FP+beta*FN+intersection+smooth)
    return Tversky_loss_inner
        


def combined_loss(alpha=0.35,beta=0.65,e_value=1e-3):
    '''
    combine binary_crossentropy and Tversky_loss
    '''
    def combined_loss_inner(y_true,y_pred):
        smooth = 1e-6
        y_true_f=K.flatten(y_true)
        y_pred_f=K.flatten(y_pred)
        intersection=K.sum(y_true_f*y_pred_f)
        tmp = tf.less(y_pred_f,0.5)

        negative_prediction_region = tf.where(tf.less(y_pred_f,0.5),1-y_pred_f,tf.zeros_like(y_pred_f))

        negative_label_region = 1-y_true_f
        FP = K.sum(negative_label_region*y_pred_f)

        FN = K.sum(negative_prediction_region*y_true_f)
        return (1-e_value) * (1-(intersection+smooth)/(alpha*FP+beta*FN+intersection+smooth)) + e_value*(K.mean(K.binary_crossentropy(y_true, y_pred), axis=-1))
   

    return combined_loss_inner


def NoduleSegAuxLoss(y_true,y_pred):
    import tensorflow as tf
    from tensorflow.python.ops import array_ops
    from keras import backend as K
    
    index = tf.where(tf.reduce_sum(y_true,axis=1)>0)
    new_y_true = tf.gather_nd(y_true,index)
    new_y_pred = tf.gather_nd(y_pred,index)
    return tf.losses.softmax_cross_entropy(new_y_true,new_y_pred)
    

def focal_loss(classes_num, gamma=2., alpha=.25, e=0.1):
    # classes_num contains sample number of each classes
    def focal_loss_fixed(target_tensor, prediction_tensor):
        '''
        prediction_tensor is the output tensor with shape [None, 100], where 100 is the number of classes
        target_tensor is the label tensor, same shape as predcition_tensor
        '''
        import tensorflow as tf
        from tensorflow.python.ops import array_ops
        from keras import backend as K

        #1# get focal loss with no balanced weight which presented in paper function (4)
        zeros = array_ops.zeros_like(prediction_tensor, dtype=prediction_tensor.dtype)
        one_minus_p = array_ops.where(tf.greater(target_tensor,zeros), target_tensor - prediction_tensor, zeros)
        FT = -1 * (one_minus_p ** gamma) * tf.log(tf.clip_by_value(prediction_tensor, 1e-8, 1.0))

        #2# get balanced weight alpha
        classes_weight = array_ops.zeros_like(prediction_tensor, dtype=prediction_tensor.dtype)

        total_num = float(sum(classes_num))
        classes_w_t1 = [ total_num / ff for ff in classes_num ]
        classes_w_tensor = tf.convert_to_tensor(classes_w_t1, dtype=prediction_tensor.dtype)
        classes_weight += classes_w_tensor

        alpha = array_ops.where(tf.greater(target_tensor, zeros), classes_weight, zeros)

        #3# get balanced focal loss
        balanced_fl = alpha * FT
        balanced_fl = tf.reduce_sum(balanced_fl)

        #4# add other op to prevent overfit
        # reference : https://spaces.ac.cn/archives/4493
        nb_classes = len(classes_num)
        fianal_loss = (1-e) * balanced_fl + e * K.categorical_crossentropy(K.ones_like(prediction_tensor)/nb_classes, prediction_tensor)

        return fianal_loss
    return focal_loss_fixed