from keras import backend as K
import tensorflow as tf
import numpy as np


def custom_dice_coefficient(y_true, y_pred, recall_weight=0.3):
    recall_weight = tf.Variable(recall_weight, dtype=tf.float32)
    regular_dice = dice_coefficient(y_true, y_pred)
    recall = lession_recall(y_true, y_pred)
    recall = tf.cast(recall, dtype=tf.float32)
    recall_addition = recall * regular_dice * recall_weight
    return regular_dice + recall_addition


def lession_recall(y_true, y_pred):
    conn_comp_true = tf.contrib.image.connected_components(tf.cast(tf.squeeze(y_true, axis=[-1]), tf.bool))
    conn_comp_pred = conn_comp_true * tf.cast(tf.squeeze(y_pred, axis=[-1]), tf.int32)

    n_conn_comp_true, _ = tf.unique(K.flatten(conn_comp_true))
    n_conn_comp_pred, _ = tf.unique(K.flatten(conn_comp_pred))
    n_conn_comp_true = tf.size(input=n_conn_comp_true) - 1
    n_conn_comp_pred = tf.size(input=n_conn_comp_pred) - 1

    recall = tf.cond(pred=tf.equal(n_conn_comp_pred, tf.Variable(0)),
                     true_fn=lambda: tf.Variable(1.0, dtype=tf.float64), false_fn=lambda: n_conn_comp_pred / n_conn_comp_true)
    return recall


def thresholded_dice(y_true, y_pred):
    y_true = tf.math.floor(y_true + 0.6)
    return dice_coefficient(y_true, y_pred)

def thresholded_dice_loss(y_true, y_pred):
    return -thresholded_dice(y_true, y_pred)

def custom_dice_coefficient_loss(y_true, y_pred):
    return -custom_dice_coefficient(y_true, y_pred)


def dice_coefficient(y_true, y_pred, smooth=0.1):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)

    intersection = K.sum(y_pred_f * y_true_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coefficient_loss(y_true, y_pred):
    return -dice_coefficient(y_true, y_pred)

def sigmoid(x):
    return 1. / (1. + K.exp(-x))

def segmentation_recall(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)

    recall = K.sum(y_pred_f * y_true_f) / tf.cast(K.sum(y_true_f), tf.float32)
    return recall


def weighted_crossentropy_pixelwise(y_true, y_pred):

    y_pred = tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7)
    y_pred = K.log(y_pred / (1 - y_pred))

    wmh_indexes = np.where(y_true == 1.0)
    weights = np.repeat(1.0, 240 * 240)
    weights = np.reshape(weights, (1, 240, 240, 1))
    weights[wmh_indexes] = 5000.0

    crossentropy = (y_true * weights * -K.log(sigmoid(y_pred)) + (1 - y_true * weights) * -K.log(1 - sigmoid(y_pred)))
    return crossentropy


def prediction_count(y_true, y_pred):
    return tf.math.count_nonzero(y_pred)


def label_count(y_true, y_pred):
    return tf.math.count_nonzero(y_true)


def prediction_sum(y_true, y_pred):
    return tf.reduce_sum(input_tensor=y_pred)


def label_sum(y_true, y_pred):
    return tf.reduce_sum(input_tensor=y_true)


custom_dice_coef = custom_dice_coefficient
custom_dice_loss = custom_dice_coefficient_loss
dice_coef = dice_coefficient
dice_coef_loss = dice_coefficient_loss

weighted_crossentropy = weighted_crossentropy_pixelwise

predicted_count = prediction_count
predicted_sum = prediction_sum

ground_truth_count = label_count
ground_truth_sum =   label_sum

pixel_recall = segmentation_recall

obj_recall = lession_recall
[((160, 204), 'tensorflow.Variable', 'tf.Variable', (['recall_weight'], {'dtype': 'tf.float32'}), '(recall_weight, dtype=tf.float32)\n', (171, 204), True, 'import tensorflow as tf\n'), ((314, 347), 'tensorflow.cast', 'tf.cast', (['recall'], {'dtype': 'tf.float32'}), '(recall, dtype=tf.float32)\n', (321, 347), True, 'import tensorflow as tf\n'), ((1197, 1224), 'tensorflow.math.floor', 'tf.math.floor', (['(y_true + 0.6)'], {}), '(y_true + 0.6)\n', (1210, 1224), True, 'import tensorflow as tf\n'), ((1528, 1545), 'keras.backend.flatten', 'K.flatten', (['y_true'], {}), '(y_true)\n', (1537, 1545), True, 'from keras import backend as K\n'), ((1561, 1578), 'keras.backend.flatten', 'K.flatten', (['y_pred'], {}), '(y_pred)\n', (1570, 1578), True, 'from keras import backend as K\n'), ((1599, 1625), 'keras.backend.sum', 'K.sum', (['(y_pred_f * y_true_f)'], {}), '(y_pred_f * y_true_f)\n', (1604, 1625), True, 'from keras import backend as K\n'), ((1910, 1927), 'keras.backend.flatten', 'K.flatten', (['y_true'], {}), '(y_true)\n', (1919, 1927), True, 'from keras import backend as K\n'), ((1943, 1960), 'keras.backend.flatten', 'K.flatten', (['y_pred'], {}), '(y_pred)\n', (1952, 1960), True, 'from keras import backend as K\n'), ((2128, 2170), 'tensorflow.clip_by_value', 'tf.clip_by_value', (['y_pred', '(1e-07)', '(1 - 1e-07)'], {}), '(y_pred, 1e-07, 1 - 1e-07)\n', (2144, 2170), True, 'import tensorflow as tf\n'), ((2182, 2210), 'keras.backend.log', 'K.log', (['(y_pred / (1 - y_pred))'], {}), '(y_pred / (1 - y_pred))\n', (2187, 2210), True, 'from keras import backend as K\n'), ((2230, 2253), 'numpy.where', 'np.where', (['(y_true == 1.0)'], {}), '(y_true == 1.0)\n', (2238, 2253), True, 'import numpy as np\n'), ((2268, 2293), 'numpy.repeat', 'np.repeat', (['(1.0)', '(240 * 240)'], {}), '(1.0, 240 * 240)\n', (2277, 2293), True, 'import numpy as np\n'), ((2308, 2345), 'numpy.reshape', 'np.reshape', (['weights', '(1, 240, 240, 1)'], {}), '(weights, (1, 240, 240, 1))\n', (2318, 2345), True, 'import numpy as np\n'), ((2575, 2604), 'tensorflow.math.count_nonzero', 'tf.math.count_nonzero', (['y_pred'], {}), '(y_pred)\n', (2596, 2604), True, 'import tensorflow as tf\n'), ((2651, 2680), 'tensorflow.math.count_nonzero', 'tf.math.count_nonzero', (['y_true'], {}), '(y_true)\n', (2672, 2680), True, 'import tensorflow as tf\n'), ((2730, 2764), 'tensorflow.reduce_sum', 'tf.reduce_sum', ([], {'input_tensor': 'y_pred'}), '(input_tensor=y_pred)\n', (2743, 2764), True, 'import tensorflow as tf\n'), ((2809, 2843), 'tensorflow.reduce_sum', 'tf.reduce_sum', ([], {'input_tensor': 'y_true'}), '(input_tensor=y_true)\n', (2822, 2843), True, 'import tensorflow as tf\n'), ((720, 745), 'keras.backend.flatten', 'K.flatten', (['conn_comp_true'], {}), '(conn_comp_true)\n', (729, 745), True, 'from keras import backend as K\n'), ((783, 808), 'keras.backend.flatten', 'K.flatten', (['conn_comp_pred'], {}), '(conn_comp_pred)\n', (792, 808), True, 'from keras import backend as K\n'), ((833, 864), 'tensorflow.size', 'tf.size', ([], {'input': 'n_conn_comp_true'}), '(input=n_conn_comp_true)\n', (840, 864), True, 'import tensorflow as tf\n'), ((892, 923), 'tensorflow.size', 'tf.size', ([], {'input': 'n_conn_comp_pred'}), '(input=n_conn_comp_pred)\n', (899, 923), True, 'import tensorflow as tf\n'), ((1975, 2001), 'keras.backend.sum', 'K.sum', (['(y_pred_f * y_true_f)'], {}), '(y_pred_f * y_true_f)\n', (1980, 2001), True, 'from keras import backend as K\n'), ((555, 584), 'tensorflow.squeeze', 'tf.squeeze', (['y_true'], {'axis': '[-1]'}), '(y_true, axis=[-1])\n', (565, 584), True, 'import tensorflow as tf\n'), ((642, 671), 'tensorflow.squeeze', 'tf.squeeze', (['y_pred'], {'axis': '[-1]'}), '(y_pred, axis=[-1])\n', (652, 671), True, 'import tensorflow as tf\n'), ((1842, 1851), 'keras.backend.exp', 'K.exp', (['(-x)'], {}), '(-x)\n', (1847, 1851), True, 'from keras import backend as K\n'), ((2012, 2027), 'keras.backend.sum', 'K.sum', (['y_true_f'], {}), '(y_true_f)\n', (2017, 2027), True, 'from keras import backend as K\n'), ((982, 996), 'tensorflow.Variable', 'tf.Variable', (['(0)'], {}), '(0)\n', (993, 996), True, 'import tensorflow as tf\n'), ((1036, 1070), 'tensorflow.Variable', 'tf.Variable', (['(1.0)'], {'dtype': 'tf.float64'}), '(1.0, dtype=tf.float64)\n', (1047, 1070), True, 'import tensorflow as tf\n'), ((1669, 1684), 'keras.backend.sum', 'K.sum', (['y_true_f'], {}), '(y_true_f)\n', (1674, 1684), True, 'from keras import backend as K\n'), ((1687, 1702), 'keras.backend.sum', 'K.sum', (['y_pred_f'], {}), '(y_pred_f)\n', (1692, 1702), True, 'from keras import backend as K\n')]