import torch.nn as nn

import torch

import numpy as np

def _fast_hist(true, pred, num_classes):

    pred = np.round(pred).astype(int)

    true = np.round(true).astype(int)

    mask = (true >= 0) & (true < num_classes)

    hist = np.bincount(

        num_classes * true[mask] + pred[mask],

        minlength=num_classes ** 2,

    ).reshape(num_classes, num_classes).astype(np.float)

    return hist

def jaccard_index(hist):

    """Computes the Jaccard index, a.k.a the Intersection over Union (IoU).

    Args:

        hist: confusion matrix.

    Returns:

        avg_jacc: the average per-class jaccard index.

    """

    A_inter_B = np.diag(hist)

    A = np.sum(hist,axis=1)

    B = np.sum(hist,axis=0)

    jaccard = A_inter_B / (A + B - A_inter_B + 1e-6)

    avg_jacc =np.nanmean(jaccard) #the mean of jaccard without NaNs

    return avg_jacc, jaccard

def dice_coef_metric(hist):

    """Computes the dice coefficient).

    Args:

        hist: confusion matrix.

     Returns:

        avg_dice: the average per-class dice coefficient.

    """

    A_inter_B = np.diag(hist)

    A = np.sum(hist,axis=1)

    B = np.sum(hist,axis=0)

    dsc = A_inter_B * 2 / (A + B + 1e-6)

    avg_dsc=np.nanmean(dsc) #the mean of dsc without NaNs

    return avg_dsc

def dice_coef_loss(y_pred, y_true):

      smooth=1.0

      assert y_pred.size() == y_true.size()

      intersection = (y_pred * y_true).sum()

      dsc = (2. * intersection + smooth) / (

          y_pred.sum() + y_true.sum() + smooth

      )

      return 1. - dsc

def bce_dice_loss(y_pred, y_true):

    dicescore = dice_coef_loss(y_pred, y_true)

    bcescore = nn.BCELoss()

    m = nn.Sigmoid()

    bceloss = bcescore(m(y_pred), y_true)

    return (bceloss + dicescore)