###########################################################################################
#                                                                                         #
# This sample shows how to evaluate object detections applying the following metrics:     #
#  * Precision x Recall curve       ---->       used by VOC PASCAL 2012)                  #
#  * Average Precision (AP)         ---->       used by VOC PASCAL 2012)                  #
#                                                                                         #
# 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 argparse
import glob
import os
import shutil
# from argparse import RawTextHelpFormatter
import sys

import _init_paths
from Label import Label
from Evaluator import *
import pdb
import json
import copy
from pprint import pprint as pp

if not (os.getenv('hikmars_version') and os.getenv('hikmars_version') == '3'):
    try:
        from hikframe.dl.base.evaluate import Evaluation
        EL = Evaluation()
    except:
        pass


def ValidateMandatoryArgs(arg, argName, errors):
    if arg is None:
        errors.append('argument %s: required argument' % argName)
    else:
        return True


def ValidatePaths(arg, nameArg, errors):
    if arg is None:
        errors.append('argument %s: invalid directory' % nameArg)
    elif os.path.isdir(arg) is False and os.path.isdir(
            os.path.join(currentPath, arg)) is False:
        errors.append('argument %s: directory does not exist \'%s\'' %
                      (nameArg, arg))
    elif os.path.isdir(os.path.join(currentPath, arg)) is True:
        arg = os.path.join(currentPath, arg)
    else:
        errors.append('argument %s: invalid directory' % nameArg)
    return arg


def getlabels(directory, isGT, allLabels=None, allClasses=None):
    """Read txt files containing bounding boxes (ground truth and detections)."""
    if allLabels is None:
        allLabels = []
    if allClasses is None:
        allClasses = []
    # Read ground truths
    os.chdir(directory)
    files = glob.glob("*.txt")

    for f in files:
        nameOfImage = f.replace(".txt", "")
        fh1 = open(f, "r")
        for line in fh1:
            line = line.replace("\n", "")
            if line.replace(' ', '') == '':
                continue
            splitLine = line.split(" ")
            if isGT:
                # idClass = int(splitLine[0]) #class
                idClass = (splitLine[0])  # class
                label = Label(nameOfImage, idClass, LabelType.GroundTruth)
            else:
                # idClass = int(splitLine[0]) #class
                idClass = (splitLine[0])  # class
                confidence = float(splitLine[1])
                label = Label(nameOfImage, idClass, LabelType.Detected,
                              confidence)

            allLabels.append(label)
            if idClass not in allClasses:
                allClasses.append(idClass)
        fh1.close()
    return allLabels, allClasses


# Get current path to set default folders
currentPath = os.path.dirname(os.path.abspath(__file__))

parser = argparse.ArgumentParser(
    prog='Object Detection Metrics - Pascal VOC',
    description='This project applies the most popular metrics used to evaluate object detection '
    'algorithms.\nThe current implemention runs the Pascal VOC metrics.\nFor further references, '
    'please check:\nhttps://github.com/rafaelpadilla/Object-Detection-Metrics',
    epilog="Developed by: Rafael Padilla (rafael.padilla@smt.ufrj.br)")

# Positional arguments
# Mandatory
parser.add_argument('-gt',
                    '--gtfolder',
                    dest='gtFolder',
                    default=os.path.join(currentPath, 'groundtruths'),
                    metavar='',
                    help='folder containing your ground truth bounding boxes')
parser.add_argument('-det',
                    '--detfolder',
                    dest='detFolder',
                    default=os.path.join(currentPath, 'detections'),
                    metavar='',
                    help='folder containing your detected bounding boxes')
# Optional
parser.add_argument('-t',
                    '--threshold',
                    dest='scoreThreshold',
                    type=float,
                    default=0.5,
                    metavar='',
                    help='score threshold. Default 0.5')
parser.add_argument('-sp',
                    '--savepath',
                    dest='savePath',
                    metavar='',
                    help='folder where the plots are saved')
parser.add_argument('-np',
                    '--noplot',
                    dest='showPlot',
                    action='store_true',
                    help='no plot is shown during execution')

parser.add_argument('-methods',
                    dest='methods',
                    metavar='',
                    help='index to be saved')

args = parser.parse_args()

scoreThreshold = args.scoreThreshold

scoreThreshold = round(scoreThreshold, 2)
print('scoreThreshold = %.2f' % scoreThreshold)

# Arguments validation
errors = []
# Groundtruth folder
if ValidateMandatoryArgs(args.gtFolder, '-gt/--gtfolder', errors):
    gtFolder = ValidatePaths(args.gtFolder, '-gt/--gtfolder', errors)
else:
    # errors.pop()
    gtFolder = os.path.join(currentPath, 'groundtruths')
    if os.path.isdir(gtFolder) is False:
        errors.append('folder %s not found' % gtFolder)

# Detection folder
if ValidateMandatoryArgs(args.detFolder, '-det/--detfolder', errors):
    detFolder = ValidatePaths(args.detFolder, '-det/--detfolder', errors)
else:
    # errors.pop()
    detFolder = os.path.join(currentPath, 'detections')
    if os.path.isdir(detFolder) is False:
        errors.append('folder %s not found' % detFolder)
if args.savePath is not None:
    savePath = ValidatePaths(args.savePath, '-sp/--savepath', errors)
else:
    savePath = os.path.join(currentPath, 'results')
# Validate savePath
# If error, show error messages
if len(errors) != 0:
    print(
        """usage: Object Detection Metrics [-h] [-v] [-gt] [-det] [-t] [-gtformat]
                                [-detformat] [-save]""")
    print('Object Detection Metrics: error(s): ')
    [print(e) for e in errors]
    sys.exit()

# generate results
results_path = os.path.join(savePath, 'results')
metrics_path = os.path.join(savePath, 'metrics')
pr_path = os.path.join(metrics_path, 'PR')
# shutil.rmtree(results_path, ignore_errors=True)
# shutil.rmtree(metrics_path, ignore_errors=True)
# os.makedirs(results_path)
os.makedirs(metrics_path, exist_ok=True)
os.makedirs(pr_path, exist_ok=True)
# # Clear folder and save results
# shutil.rmtree(savePath, ignore_errors=True)
# os.makedirs(savePath)
# Show plot during execution
showPlot = args.showPlot

methods_show = args.methods
methods_show = methods_show.split(',')
print('methods:', methods_show)

# Get groundtruth boxes
allLabels, allClasses = getlabels(gtFolder, True)
# Get detected boxes
allLabels, allClasses = getlabels(detFolder, False, allLabels, allClasses)
allClasses.sort()

evaluator = Evaluator()
acc_AP = 0

macro_prec = 0
macro_recall = 0
macro_f1score = 0

all_tp = 0
all_sample = 0

# metrics.json
metric_dict = {
    "index": {
        "brief": {},
        "detail": {}
    },
    "chart": {
        "FP_view": {
            "right": 0,
            "wrong": 0
        },
        "FN_view": {
            "right": 0,
            "wrong": 0
        }
    }
}
"""
Acc,Macro-Precision,Macro-Recall,Macro-F1-Score
"""
detections = evaluator.NormalMetric(
    results_path,
    # Object containing all bounding boxes (ground truths and detections)
    allLabels,
    scoreThreshold)

# f = open(metric_path, 'a')
# f.write('Precision/Recall/f1score/acc per class(scoreThreshold={}):'.format(scoreThreshold))
# each detection is a class
validClasses = 0
for metricsPerClass in detections:
    # Get metric values per each class
    cl = metricsPerClass['class']
    precision = metricsPerClass['precision']
    recall = metricsPerClass['recall']
    f1score = metricsPerClass['f1score']
    totalTP = metricsPerClass['total TP']
    totalPositives = metricsPerClass['total positives']

    macro_prec += precision
    macro_recall += recall
    macro_f1score += f1score

    if totalPositives > 0:
        validClasses = validClasses + 1
        p_str = "{0:.2f}%".format(precision * 100)
        r_str = "{0:.2f}%".format(recall * 100)
        f_str = "{0:.2f}%".format(f1score * 100)
        print('Precision: {} ({})'.format(p_str, cl))
        print('Recall: {} ({})'.format(r_str, cl))
        print('F1-Score: {} ({})'.format(f_str, cl))
        metric_dict["index"]["detail"][cl] = {}
        cls_detail = metric_dict["index"]["detail"][cl]
        if "Macro-Precision" in methods_show:
            cls_detail["Precision"] = p_str
        if "Macro-Recall" in methods_show:
            cls_detail["Recall"] = r_str
        if "Macro-F1-Score" in methods_show:
            cls_detail["F1-Score"] = f_str

        all_tp += totalTP
        all_sample += totalPositives

macro_prec = macro_prec / validClasses
macro_recall = macro_recall / validClasses
macro_f1score = macro_f1score / validClasses

macro_prec_str = "{0:.2f}%".format(macro_prec * 100)
print('Macro-Precision: {}'.format(macro_prec_str))
# f.write('\nMacro-Precision: %s\n' % macro_prec_str)
macro_recall_str = "{0:.2f}%".format(macro_recall * 100)
print('Macro-Recall: {}'.format(macro_recall_str))
# f.write('\nMacro-Recall: %s\n' % macro_recall_str)
macro_f1score_str = "{0:.2f}%".format(macro_f1score * 100)
print('Macro-F1-Score: {}'.format(macro_f1score_str))
# f.write('\nMacro-F1-Score: %s\n' % macro_f1score_str)

acc = all_tp / all_sample
acc_str = "{0:.2f}%".format(acc * 100)
print('Accuracy: {}'.format(acc_str))
# f.write('\nAcc: %s\n' % acc_str)
# f.close()
cls_brief = metric_dict["index"]["brief"]
if "Macro-Precision" in methods_show:
    cls_brief["Macro-Precision"] = macro_prec_str
if "Macro-Recall" in methods_show:
    cls_brief["Macro-Recall"] = macro_recall_str
if "Macro-F1-Score" in methods_show:
    cls_brief["Macro-F1-Score"] = macro_f1score_str
if "Accuracy" in methods_show:
    cls_brief["Accuracy"] = acc_str

FP_view = metric_dict["chart"]["FP_view"]
FP_view["right"] = int(all_tp)
FP_view["wrong"] = int(all_sample - all_tp)
metric_dict["chart"]["FN_view"] = copy.copy(FP_view)
print('---')
"""
mAP
"""
# Plot Precision x Recall curve
detections = evaluator.PlotPrecisionRecallCurve(
    # Object containing all bounding boxes (ground truths and detections)
    allLabels,
    method=MethodAveragePrecision.EveryPointInterpolation,
    showAP=True,  # Show Average Precision in the title of the plot
    showInterpolatedPrecision=False,  # Don't plot the interpolated precision curve
    savePath=pr_path,
    showGraphic=showPlot)

# f = open(metric_path, 'a')
# f.write('Average Precision (AP), Precision and Recall per class:')

# each detection is a class

validClasses = 0
for metricsPerClass in detections:
    # Get metric values per each class
    cl = metricsPerClass['class']
    ap = metricsPerClass['AP']
    # precision = metricsPerClass['precision']
    # recall = metricsPerClass['recall']
    # f1score = metricsPerClass['f1score']
    totalPositives = metricsPerClass['total positives']

    if totalPositives > 0:
        validClasses = validClasses + 1
        acc_AP = acc_AP + ap
        # prec = ['%.2f' % p for p in precision]
        # rec = ['%.2f' % r for r in recall]
        ap_str = "{0:.2f}%".format(ap * 100)
        # ap_str = "{0:.4f}%".format(ap * 100)
        cls_detail = metric_dict["index"]["detail"][cl]
        if "mAP" in methods_show:
            print('AP: %s (%s)' % (ap_str, cl))
            cls_detail["AP"] = ap_str
        pr_img_path = os.path.join(pr_path, '{}.png'.format(cl))
        if os.path.exists(pr_img_path):
            cls_detail["pr_img_path"] = pr_img_path

        # all_sample += totalPositives

mAP = acc_AP / validClasses
mAP_str = "{0:.2f}%".format(mAP * 100)
if "mAP" in methods_show:
    print('mAP: %s' % mAP_str)
cls_brief = metric_dict["index"]["brief"]
if "mAP" in methods_show:
    cls_brief["mAP"] = mAP_str

# tmp_list = []
# for k, v in metric_dict["index"]["detail"].items():
#     tmp_list.append({k: v})
# metric_dict["index"]["detail"] = tmp_list
pp(metric_dict)


# modify metric_dict
f_metric_dict = {
    "index": {
        "brief": {},
        "detail": {}
    },
    "chart": {
        "误检性分析": {
            "正确分类": 0,
            "错误分类": 0,
        }
    }
}


f_metric_dict["index"] = metric_dict["index"]
f_metric_dict["chart"]["误检性分析"] = {
    "正确分类": metric_dict["chart"]["FP_view"]["right"],
    "错误分类": metric_dict["chart"]["FP_view"]["wrong"]}
metric_file = os.path.join(metrics_path, 'metrics.json')
shutil.rmtree(metric_file, ignore_errors=True)


# fit hikmars3
if os.getenv('hikmars_version') and os.getenv('hikmars_version') == '3':
    with open(os.path.join(metric_file), "w") as f:
        json.dump(f_metric_dict, f)
else:
    brief = f_metric_dict["index"]["brief"]
    index_names = []
    index_values = []
    for k, v in brief.items():
        index_names.append(k)
        index_values.append(v)
    EL.brief_index(index_names, index_values)

    detail = f_metric_dict["index"]["detail"]
    for k, v in detail.items():
        label_name = k
        data = v
        index_names = []
        index_values = []
        for name, value in data.items():
            if name == 'pr_img_path':
                continue
            index_names.append(name)
            index_values.append(value)
        image_path = data['pr_img_path']
        EL.detail_index(label_name, index_names, index_values, image_path)

    for k, v in f_metric_dict["chart"].items():
        name = k
        data = []
        for name, value in v.items():
            data.append({"name": name, "value": int(value)})
        mode = 'bar'
        EL.chart(name, data, mode)
