"""Evaluation Metrics for Semantic Segmentation"""
import numpy as np
import mxnet.ndarray as F

__all__ = ['batch_pix_accuracy', 'batch_intersection_union', 'pixelAccuracy',
           'intersectionAndUnion']

def batch_pix_accuracy(output, target, ignore_bg=False):
    """PixAcc"""
    # inputs are NDarray, output 4D, target 3D
    predict = F.argmax(output, 1)
    target = target.astype(predict.dtype)
    if ignore_bg:
        pixel_labeled = (target > 0).sum().asscalar()
        pixel_correct = (F.equal(predict, target)*(target > 0.0)).sum().asscalar()
    else:
        pixel_labeled = (target >= 0).sum().asscalar()
        pixel_correct = (F.equal(predict, target)*(target >= 0.0)).sum().asscalar()
    return pixel_correct, pixel_labeled


def batch_intersection_union(output, target, nclass, ignore_bg=False):
    """mIoU"""
    # inputs are NDarray, output 4D, target 3D
    # ignore_bg=True, ignoring class 0; ignore_bg = False, use class 0
    predict = F.argmax(output, 1)
    target = target.astype(predict.dtype)
    mini = 0
    maxi = nclass - 1
    nbins = nclass
    predict = predict.asnumpy()
    target = target.asnumpy()
    if ignore_bg:
        mini = 1
        nbins -= 1
        predict = predict * (target > 0).astype(predict.dtype)
    else:
        predict = predict * (target >= 0).astype(predict.dtype)
    #intersection = predict * (F.equal(predict, target)).astype(predict.dtype)
    intersection = predict * (predict == target)
    # areas of intersection and union
    area_inter, _ = np.histogram(intersection, bins=nbins,
                                 range=(mini, maxi))
    area_pred, _ = np.histogram(predict, bins=nbins,
                                range=(mini, maxi))
    area_lab, _ = np.histogram(target, bins=nbins,
                               range=(mini, maxi))
    area_union = area_pred + area_lab - area_inter
    return area_inter, area_union


def pixelAccuracy(imPred, imLab):
    """
    This function takes the prediction and label of a single image, returns pixel-wise accuracy
    To compute over many images do:
    for i = range(Nimages):
         (pixel_accuracy[i], pixel_correct[i], pixel_labeled[i]) = \
            pixelAccuracy(imPred[i], imLab[i])
    mean_pixel_accuracy = 1.0 * np.sum(pixel_correct) / (np.spacing(1) + np.sum(pixel_labeled))
    """
    # Remove classes from unlabeled pixels in gt image.
    # We should not penalize detections in unlabeled portions of the image.
    pixel_labeled = np.sum(imLab > 0)
    pixel_correct = np.sum((imPred == imLab)*(imLab > 0))
    pixel_accuracy = 1.0 * pixel_correct / pixel_labeled
    return (pixel_accuracy, pixel_correct, pixel_labeled)


def intersectionAndUnion(imPred, imLab, numClass):
    """
    This function takes the prediction and label of a single image,
    returns intersection and union areas for each class
    To compute over many images do:
    for i in range(Nimages):
        (area_intersection[:,i], area_union[:,i]) = intersectionAndUnion(imPred[i], imLab[i])
    IoU = 1.0 * np.sum(area_intersection, axis=1) / np.sum(np.spacing(1)+area_union, axis=1)
    """
    # Remove classes from unlabeled pixels in gt image.
    # We should not penalize detections in unlabeled portions of the image.
    imPred = imPred * (imLab > 0)

    # Compute area intersection:
    intersection = imPred * (imPred == imLab)
    (area_intersection, _) = np.histogram(intersection, bins=numClass, range=(1, numClass))

    # Compute area union:
    (area_pred, _) = np.histogram(imPred, bins=numClass, range=(1, numClass))
    (area_lab, _) = np.histogram(imLab, bins=numClass, range=(1, numClass))
    area_union = area_pred + area_lab - area_intersection
    return (area_intersection, area_union)
