import numpy as np

def single_dice_coef(y_pred, y_true):

    # shape of y_true and y_pred: (height, width)

    intersection = np.sum(y_true * y_pred)

    if (np.sum(y_true) == 0) and (np.sum(y_pred) == 0):

        return 1

    return (2*intersection) / (np.sum(y_true) + np.sum(y_pred))

def mean_dice_coef(y_pred, y_true):

    # shape of y_true and y_pred: (n_samples, height, width)

    batch_size = y_true.shape[0]

    mean_dice_channel = 0.

    for i in range(batch_size):

        channel_dice = single_dice_coef(y_pred[i, :, :], y_true[i, :, :])

        mean_dice_channel += channel_dice/(batch_size)

    return mean_dice_channel

def mean_dice_coef_remove_empty(y_pred, y_true):

    # shape of y_true and y_pred: (n_samples, height, width)

    batch_size = y_true.shape[0]

    mean_dice_channel = 0.

    num_no_empty = batch_size

    for i in range(batch_size):

        if (np.sum(y_true[i, :, :]) == 0):

            num_no_empty -= 1

            continue

        channel_dice = single_dice_coef(y_pred[i, :, :], y_true[i, :, :])

        mean_dice_channel += channel_dice

    if num_no_empty == 0:

        return None

    return mean_dice_channel/(num_no_empty)