from keras import backend as K

from keras.losses import binary_crossentropy





#########################################*Functions*####################################################################
#-----------------------------------------------------------------------------------------------------------------------
def dice_coef_old(y_true, y_pred):

	y_true_f = K.flatten(y_true)
	y_pred = K.cast(y_pred, 'float32')
	y_pred_f = K.cast(K.greater(K.flatten(y_pred), 0.5), 'float32')
	intersection = y_true_f * y_pred_f
	score = 2.0 * K.sum(intersection) / (K.sum(y_true_f) + K.sum(y_pred_f))

	return score
#-----------------------------------------------------------------------------------------------------------------------
def dice_coef(y_true, y_pred, smooth=1.):

	intersection = K.sum(y_true * y_pred, axis=[1,2,3])
	union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3])

	return K.mean((2. * intersection + smooth) / (union + smooth), axis=0)
#-----------------------------------------------------------------------------------------------------------------------
def dice_coef_log(y_true, y_pred, smooth=1.):

	intersection = K.sum(y_true * y_pred, axis=[1,2,3])
	union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3])

	return 1. + K.log(K.mean((2. * intersection + smooth) / (union + smooth), axis=0))

#-----------------------------------------------------------------------------------------------------------------------
def dice_loss(y_true, y_pred):

	smooth = 1.0
	y_true_f = K.flatten(y_true)
	y_pred_f = K.flatten(y_pred)
	intersection = y_true_f * y_pred_f
	score = (2. * K.sum(intersection) + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

	return 1.0 - score
#-----------------------------------------------------------------------------------------------------------------------
def bce_dice_loss(y_true, y_pred):

	return binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)
#-----------------------------------------------------------------------------------------------------------------------
def bce_logdice_loss(y_true, y_pred):

	return binary_crossentropy(y_true, y_pred) - K.log(1.0 - dice_loss(y_true, y_pred))
#-----------------------------------------------------------------------------------------------------------------------
def weighted_bce_loss(y_true, y_pred, weight):

	epsilon = 1e-7
	y_pred = K.clip(y_pred, epsilon, 1.0 - epsilon)
	logit_y_pred = K.log(y_pred / (1.0 - y_pred))
	loss = weight * (logit_y_pred * (1.0 - y_true) + K.log(1. + K.exp(-K.abs(logit_y_pred))) +
	                 K.maximum(-logit_y_pred, 0.))

	return K.sum(loss) / K.sum(weight)
#-----------------------------------------------------------------------------------------------------------------------
def weighted_dice_loss(y_true, y_pred, weight):

	smooth = 1.0
	w, m1, m2 = weight, y_true, y_pred
	intersection = (m1 * m2)
	score = (2.0 * K.sum(w * intersection) + smooth) / (K.sum(w * m1) + K.sum(w * m2) + smooth)
	loss = 1.0 - K.sum(score)

	return loss
#-----------------------------------------------------------------------------------------------------------------------
def weighted_bce_dice_loss(y_true, y_pred):

	y_true = K.cast(y_true, 'float32')
	y_pred = K.cast(y_pred, 'float32')
	# if we want to get same size of output, kernel size must be odd
	averaged_mask = K.pool2d(y_true, pool_size=(50, 50), strides=(1, 1), padding='same', pool_mode='avg')
	weight = K.ones_like(averaged_mask)
	w0 = K.sum(weight)
	weight = 5. * K.exp(-5. * K.abs(averaged_mask - 0.5))
	w1 = K.sum(weight)
	weight *= (w0 / w1)
	loss = weighted_bce_loss(y_true, y_pred, weight) + dice_loss(y_true, y_pred)

	return loss
#-----------------------------------------------------------------------------------------------------------------------
def dice_p_bce(in_gt, in_pred):

	return 1e-3 * binary_crossentropy(in_gt, in_pred) - dice_coef(in_gt, in_pred)
#-----------------------------------------------------------------------------------------------------------------------
def true_positive_rate(y_true, y_pred):

	return K.sum(K.flatten(y_true) * K.flatten(K.round(y_pred))) / (K.sum(y_true) + K.epsilon())
#-----------------------------------------------------------------------------------------------------------------------