###########################################################################################
#                                                                                         #
# Evaluator class: Implements the most popular metrics for object detection               #
#                                                                                         #
# Developed by: Rafael Padilla (rafael.padilla@smt.ufrj.br)                               #
#        SMT - Signal Multimedia and Telecommunications Lab                               #
#        COPPE - Universidade Federal do Rio de Janeiro                                   #
#        Last modification: Oct 9th 2018                                                 #
###########################################################################################

import pdb
from utils import *
from Label import *
import numpy as np
import os
import sys
from collections import Counter
import pprint
import shutil
import matplotlib.pyplot as plt
plt.switch_backend('agg')

os.environ['debug_mode'] = 'True'


def debug_print(*args):
    if 'debug_mode' in os.environ:
        print(*args)


class Evaluator:
    def GetMetrics(self, results_path, labels, scoreThreshold=0.5):
        right_path = os.path.join(results_path, 'right')
        wrong_path = os.path.join(results_path, 'wrong')
        os.makedirs(right_path, exist_ok=True)
        os.makedirs(wrong_path, exist_ok=True)
        # list containing metrics (precision, recall, average precision) of each class
        ret = []
        groundTruths = []
        detections = []
        classes = []
        for label in labels:
            if label.getLabelType() == LabelType.GroundTruth:
                groundTruths.append(
                    [label.getImageName(),
                     label.getClassId(), 1])
                # get class
                if label.getClassId() not in classes:
                    classes.append(label.getClassId())
            else:
                detections.append([
                    label.getImageName(),
                    label.getClassId(),
                    label.getConfidence()
                ])

        classes = sorted(classes)
        # Precision x Recall is obtained individually by each class
        # Loop through by classes

        for c in classes:
            # Get only detection of class c
            dects = []
            [
                # dects.append(d) for d in detections if (d[1] == c and d[2] >= scoreThreshold)
                dects.append(d) for d in detections if (d[1] == c)
            ]
            # Get only ground truths of class c
            gts = []
            [gts.append(g) for g in groundTruths if g[1] == c]

            npos = len(gts)
            # sort detections by decreasing confidence
            dects = sorted(dects, key=lambda conf: conf[2], reverse=True)
            TP = np.zeros(len(dects))
            FP = np.zeros(len(dects))
            gt_imgs = [cc[0] for cc in gts]

            for d in range(len(dects)):
                img_nm = dects[d][0]
                if dects[d][0] in gt_imgs:
                    TP[d] = 1  # count as true positive
                    # shutil.move(
                    #     '{}/{}.json'.format(results_path, img_nm), right_path)
                else:
                    FP[d] = 1  # count as false positive
                    json_path = '{}/{}.json'.format(right_path, img_nm)
                    try:
                        shutil.move(json_path, wrong_path)
                    except:
                        print('no such file {}, skipped.'.format(json_path))
                        pass
            # compute precision, recall and average precision
            acc_TP = np.cumsum(TP)
            acc_FP = np.cumsum(FP)

            # rec = (acc_TP / npos) if npos > 0 else 0
            # prec = acc_TP / (acc_FP + acc_TP) if (acc_FP + acc_TP) > 0 else 0
            # f1score = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0
            # debug_print(prec)
            # debug_print(rec)
            # debug_print(f1score)

            # cal pre and recall
            confs = [x[2] for x in dects]
            rec = acc_TP / npos if npos > 0 else np.zeros_like(acc_TP)
            # rec = acc_TP / npos
            prec = np.divide(acc_TP, (acc_FP + acc_TP))

            debug_print('[!!!]precision and recall calculate detail:')
            debug_print('     all det_box confs(decrease):', confs)
            debug_print('     all det_box precsions_related(decrease):', prec)
            debug_print('     all det_box recalls_related(increase):', rec)
            ret_prec = 0
            ret_rec = 0
            flag = 0

            if confs:
                debug_print('     reason: scoreThreshold', scoreThreshold)
                if confs[0] < scoreThreshold:
                    debug_print('     reason: confs[0]', confs[0])
                    debug_print('     reason: confs[0] < scoreThreshold')
                    ret_prec = 0
                    ret_rec = 0
                else:
                    for i, conf in enumerate(confs):
                        if conf < scoreThreshold:
                            # ret_c = confs[i-1]
                            ret_prec = prec[i-1]
                            ret_rec = rec[i-1]
                            flag = 1
                            debug_print(
                                '     reason: confs[{}] ='.format(i-1), confs[i-1])
                            debug_print(
                                '     reason: confs[{}] >= scoreThreshold'.format(i-1))
                            debug_print(
                                '     reason: confs[{}] ='.format(i), confs[i])
                            debug_print(
                                '     reason: confs[{}] < scoreThreshold'.format(i))
                            debug_print(
                                '     ret: use confs[{}]\'s conf({})'.format(i-1, confs[i-1]))
                            debug_print(
                                '     ret: get its precision and recall')
                            break
                    if not flag:
                        debug_print('     reason: no conf < scoreThreshold')
                        debug_print(
                            '     ret: use final conf({})'.format(confs[-1]))
                        debug_print('     ret: get its precision and recall')
                        ret_prec = prec[-1]
                        ret_rec = rec[-1]
            else:
                debug_print('     reason: no det_box')
            debug_print('     ret: precision=', ret_prec)
            debug_print('     ret: recall=', ret_rec)

            # f1score = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0
            f1score = 2 * ret_prec * ret_rec / \
                (ret_prec + ret_rec) if (ret_prec + ret_rec) > 0 else 0
            debug_print('     ret: f1score=', f1score)

            # add class result in the dictionary to be returned
            r = {
                'scoreThreshold': scoreThreshold,
                'class': c,
                'precision': ret_prec,
                'recall': ret_rec,
                'f1score': f1score,
                'total TP': np.sum(TP),
                'total FP': np.sum(FP),
                'total positives': npos
            }
            ret.append(r)
        # pdb.set_trace()
        pprint.pprint(ret)
        return ret

    def GetPascalVOCMetrics(self,
                            labels,
                            method=MethodAveragePrecision.
                            EveryPointInterpolation):
        """Get the metrics used by the VOC Pascal 2012 challenge.
        Get[]
        Args:
            labels: Object of the class BoundingBoxes representing ground truth and detected
            bounding boxes;
            scoreThreshold: IOU threshold indicating which detections will be considered TP or FP
            (default value = 0.5);
            method (default = EveryPointInterpolation): It can be calculated as the implementation
            in the official PASCAL VOC toolkit (EveryPointInterpolation), or applying the 11-point
            interpolatio as described in the paper "The PASCAL Visual Object Classes(VOC) Challenge"
            or EveryPointInterpolation"  (ElevenPointInterpolation);
        Returns:
            A list of dictionaries. Each dictionary contains information and metrics of each class.
            The keys of each dictionary are:
            dict['class']: class representing the current dictionary;
            dict['precision']: array with the precision values;
            dict['recall']: array with the recall values;
            dict['AP']: average precision;
            dict['interpolated precision']: interpolated precision values;
            dict['interpolated recall']: interpolated recall values;
            dict['total positives']: total number of ground truth positives;
            dict['total TP']: total number of True Positive detections;
            dict['total FP']: total number of False Positive detections;
        """
        ret = [
        ]  # list containing metrics (precision, recall, average precision) of each class
        # List with all ground truths (Ex: [imageName,class,confidence=1, (label coordinates XYX2Y2)])
        groundTruths = []
        # List with all detections (Ex: [imageName,class,confidence,(label coordinates XYX2Y2)])
        detections = []
        # Get all classes
        classes = []
        # Loop through all bounding boxes and separate them into GTs and detections
        for label in labels:
            # [imageName, class, confidence, (label coordinates XYX2Y2)]
            if label.getLabelType() == LabelType.GroundTruth:
                groundTruths.append(
                    [label.getImageName(),
                     label.getClassId(), 1])
            else:
                detections.append([
                    label.getImageName(),
                    label.getClassId(),
                    label.getConfidence()
                ])
            # get class
            if label.getClassId() not in classes:
                classes.append(label.getClassId())
        classes = sorted(classes)
        # Precision x Recall is obtained individually by each class
        # Loop through by classes

        for c in classes:
            # Get only detection of class c
            dects = []
            [dects.append(d) for d in detections if d[1] == c]
            # Get only ground truths of class c
            gts = []
            [gts.append(g) for g in groundTruths if g[1] == c]
            npos = len(gts)
            # sort detections by decreasing confidence
            dects = sorted(dects, key=lambda conf: conf[2], reverse=True)
            TP = np.zeros(len(dects))
            FP = np.zeros(len(dects))
            gt_imgs = [cc[0] for cc in gts]

            for d in range(len(dects)):
                if dects[d][0] in gt_imgs:
                    TP[d] = 1  # count as true positive
                else:
                    FP[d] = 1  # count as false positive
            # compute precision, recall and average precision
            acc_FP = np.cumsum(FP)
            acc_TP = np.cumsum(TP)

            # rec = acc_TP / npos
            # rec = (acc_TP / npos) if npos > 0 else 0
            rec = (acc_TP / npos) if npos > 0 else np.zeros_like(acc_TP)

            prec = np.divide(acc_TP, (acc_FP + acc_TP))
            # f1score = np.nan_to_num(2 * prec * rec / (prec + rec))
            # import pdb
            # pdb.set_trace()
            # Depending on the method, call the right implementation
            if method == MethodAveragePrecision.EveryPointInterpolation:
                [ap, mpre, mrec,
                 ii] = Evaluator.CalculateAveragePrecision(rec, prec)
            else:
                [ap, mpre, mrec,
                 _] = Evaluator.ElevenPointInterpolatedAP(rec, prec)
            # add class result in the dictionary to be returned
            r = {
                'class': c,
                'precision': prec,
                'recall': rec,
                # 'f1score': f1score,
                'AP': ap,
                'interpolated precision': mpre,
                'interpolated recall': mrec,
                'total positives': npos,
                'total TP': np.sum(TP),
                'total FP': np.sum(FP)
            }
            ret.append(r)
            # pdb.set_trace()
        return ret

    def NormalMetric(self, results_path, labels, scoreThreshold=0.5):
        results = self.GetMetrics(results_path, labels, scoreThreshold)
        return results

    def PlotPrecisionRecallCurve(
            self,
            labels,
            method=MethodAveragePrecision.EveryPointInterpolation,
            showAP=False,
            showInterpolatedPrecision=False,
            savePath=None,
            showGraphic=True):
        """PlotPrecisionRecallCurve
        Plot the Precision x Recall curve for a given class.
        Args:
            labels: Object of the class BoundingBoxes representing ground truth and detected
            bounding boxes;
            scoreThreshold (optional): IOU threshold indicating which detections will be considered
            TP or FP (default value = 0.5);
            method (default = EveryPointInterpolation): It can be calculated as the implementation
            in the official PASCAL VOC toolkit (EveryPointInterpolation), or applying the 11-point
            interpolatio as described in the paper "The PASCAL Visual Object Classes(VOC) Challenge"
            or EveryPointInterpolation"  (ElevenPointInterpolation).
            showAP (optional): if True, the average precision value will be shown in the title of
            the graph (default = False);
            showInterpolatedPrecision (optional): if True, it will show in the plot the interpolated
             precision (default = False);
            savePath (optional): if informed, the plot will be saved as an image in this path
            (ex: /home/mywork/ap.png) (default = None);
            showGraphic (optional): if True, the plot will be shown (default = True)
        Returns:
            A list of dictionaries. Each dictionary contains information and metrics of each class.
            The keys of each dictionary are:
            dict['class']: class representing the current dictionary;
            dict['precision']: array with the precision values;
            dict['recall']: array with the recall values;
            dict['AP']: average precision;
            dict['interpolated precision']: interpolated precision values;
            dict['interpolated recall']: interpolated recall values;
            dict['total positives']: total number of ground truth positives;
            dict['total TP']: total number of True Positive detections;
            dict['total FP']: total number of False Negative detections;
        """
        results = self.GetPascalVOCMetrics(labels, method)
        # results = self.GetPascalVOCMetrics(labels, scoreThreshold, method)
        result = None
        # Each resut represents a class
        for result in results:
            if result is None:
                raise IOError('Error: Class %d could not be found.' % classId)

            classId = result['class']
            precision = result['precision']
            recall = result['recall']
            # f1score = result['f1score']
            average_precision = result['AP']
            mpre = result['interpolated precision']
            mrec = result['interpolated recall']
            npos = result['total positives']
            total_tp = result['total TP']
            total_fp = result['total FP']

            plt.close()
            if showInterpolatedPrecision:
                if method == MethodAveragePrecision.EveryPointInterpolation:
                    plt.plot(mrec,
                             mpre,
                             '--r',
                             label='Interpolated precision (every point)')
                elif method == MethodAveragePrecision.ElevenPointInterpolation:
                    # Uncomment the line below if you want to plot the area
                    # plt.plot(mrec, mpre, 'or', label='11-point interpolated precision')
                    # Remove duplicates, getting only the highest precision of each recall value
                    nrec = []
                    nprec = []
                    for idx in range(len(mrec)):
                        r = mrec[idx]
                        if r not in nrec:
                            idxEq = np.argwhere(mrec == r)
                            nrec.append(r)
                            nprec.append(max([mpre[int(id)] for id in idxEq]))
                    plt.plot(nrec,
                             nprec,
                             'or',
                             label='11-point interpolated precision')
            plt.plot(recall, precision, label='Precision')
            plt.xlabel('recall')
            plt.ylabel('precision')
            if showAP:
                ap_str = "{0:.2f}%".format(average_precision * 100)
                # ap_str = "{0:.4f}%".format(average_precision * 100)
                plt.title('Precision x Recall curve \nClass: %s, AP: %s' %
                          (str(classId), ap_str))
            else:
                plt.title('Precision x Recall curve \nClass: %s' %
                          str(classId))
            plt.legend(shadow=True)
            plt.grid()

            if savePath is not None:
                plt.savefig(os.path.join(savePath, classId + '.png'))
            if showGraphic is True:
                plt.show()
                # plt.waitforbuttonpress()
                plt.pause(0.05)
        return results

    @ staticmethod
    def CalculateAveragePrecision(rec, prec):
        mrec = []
        mrec.append(0)
        [mrec.append(e) for e in rec]
        mrec.append(1)
        mpre = []
        mpre.append(0)
        [mpre.append(e) for e in prec]
        mpre.append(0)
        for i in range(len(mpre) - 1, 0, -1):
            mpre[i - 1] = max(mpre[i - 1], mpre[i])
        ii = []
        for i in range(len(mrec) - 1):
            if mrec[1:][i] != mrec[0:-1][i]:
                ii.append(i + 1)
        ap = 0
        for i in ii:
            ap = ap + np.sum((mrec[i] - mrec[i - 1]) * mpre[i])
        # return [ap, mpre[1:len(mpre)-1], mrec[1:len(mpre)-1], ii]
        return [ap, mpre[0:len(mpre) - 1], mrec[0:len(mpre) - 1], ii]

    @ staticmethod
    # 11-point interpolated average precision
    def ElevenPointInterpolatedAP(rec, prec):
        # def CalculateAveragePrecision2(rec, prec):
        mrec = []
        # mrec.append(0)
        [mrec.append(e) for e in rec]
        # mrec.append(1)
        mpre = []
        # mpre.append(0)
        [mpre.append(e) for e in prec]
        # mpre.append(0)
        recallValues = np.linspace(0, 1, 11)
        recallValues = list(recallValues[::-1])
        rhoInterp = []
        recallValid = []
        # For each recallValues (0, 0.1, 0.2, ... , 1)
        for r in recallValues:
            # Obtain all recall values higher or equal than r
            argGreaterRecalls = np.argwhere(mrec[:] >= r)
            pmax = 0
            # If there are recalls above r
            if argGreaterRecalls.size != 0:
                pmax = max(mpre[argGreaterRecalls.min():])
            recallValid.append(r)
            rhoInterp.append(pmax)
        # By definition AP = sum(max(precision whose recall is above r))/11
        ap = sum(rhoInterp) / 11
        # Generating values for the plot
        rvals = []
        rvals.append(recallValid[0])
        [rvals.append(e) for e in recallValid]
        rvals.append(0)
        pvals = []
        pvals.append(0)
        [pvals.append(e) for e in rhoInterp]
        pvals.append(0)
        # rhoInterp = rhoInterp[::-1]
        cc = []
        for i in range(len(rvals)):
            p = (rvals[i], pvals[i - 1])
            if p not in cc:
                cc.append(p)
            p = (rvals[i], pvals[i])
            if p not in cc:
                cc.append(p)
        recallValues = [i[0] for i in cc]
        rhoInterp = [i[1] for i in cc]
        return [ap, rhoInterp, recallValues, None]
