import tensorflow.keras.backend as K


def class_tversky(y_true, y_pred):
    smooth = 1

    y_true = K.permute_dimensions(y_true, (3, 1, 2, 0))
    y_pred = K.permute_dimensions(y_pred, (3, 1, 2, 0))

    y_true_pos = K.batch_flatten(y_true)
    y_pred_pos = K.batch_flatten(y_pred)
    true_pos = K.sum(y_true_pos * y_pred_pos, 1)
    false_neg = K.sum(y_true_pos * (1-y_pred_pos), 1)
    false_pos = K.sum((1-y_true_pos)*y_pred_pos, 1)
    alpha = 0.7
    return (true_pos + smooth)/(true_pos + alpha*false_neg + (1-alpha)*false_pos + smooth)


def focal_tversky_loss(y_true, y_pred):
    pt_1 = class_tversky(y_true, y_pred)
    gamma = 0.75
    focal_tv = K.sum(K.pow((1-pt_1), gamma))
    return focal_tv


# Ref: salehi17, "Twersky loss function for image segmentation using 3D FCDN"
# -> the score is computed for each class separately and then summed
# alpha=beta=0.5 : dice coefficient
# alpha=beta=1   : tanimoto coefficient (also known as jaccard)
# alpha+beta=1   : produces set of F*-scores
# implemented by E. Moebel, 06/04/18 and modified by me
def tversky_loss(y_true, y_pred):
    alpha = 0.7
    beta = 0.3

    ones = K.ones_like(y_true)
    p0 = y_pred      # proba that voxels are class i
    p1 = ones-y_pred  # proba that voxels are not class i
    g0 = y_true
    g1 = ones-y_true

    num = K.sum(p0*g0, (0, 1, 2))
    den = num + alpha*K.sum(p0*g1, (0, 1, 2)) + \
        beta*K.sum(p1*g0, (0, 1, 2))

    # when summing over classes, T has dynamic range [0 Ncl]
    T = K.sum(num/den)

    Ncl = K.cast(K.shape(y_true)[-1], 'float32')
    return (Ncl-T) / Ncl


def focal_tv_loss(y_true, y_pred):
    gamma = 0.75
    tl = tversky_loss(y_true, y_pred)
    ftl = K.pow(1 - tl, gamma)
    return ftl
