import torch


def compute_ap(recall, precision):
    """ Compute the average precision, given the recall and precision curves.
    Code originally from https://github.com/rbgirshick/py-faster-rcnn.

    # Arguments
        recall:    The recall curve (list).
        precision: The precision curve (list).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """
    # correct AP calculation
    # first append sentinel values at the end
    mrec = torch.cat((torch.zeros((1, ), device=recall.device, dtype=recall.dtype),
                      recall,
                      torch.ones((1, ), device=recall.device, dtype=recall.dtype)))
    mpre = torch.cat((torch.zeros((1, ), device=precision.device, dtype=precision.dtype),
                      precision,
                      torch.zeros((1, ), device=precision.device, dtype=precision.dtype)))

    # compute the precision envelope
    for i in range(len(mpre) - 1, 0, -1):
        mpre[i - 1] = torch.max(mpre[i - 1], mpre[i])

    # to calculate area under PR curve, look for points
    # where X axis (recall) changes value
    i = torch.nonzero(mrec[1:] != mrec[:-1])

    # and sum (\Delta recall) * prec
    ap = torch.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap