import math
import pickle
import numpy as np
import os 
import copy
from yolox_microbt.data.evaluation import utils


class MAPEvaluator(object):
    def __init__(self,
                 label_names,
                 iou_threshs,
                 ignore_thresh=0.5,
                 ignore_mode='ioa',
                 sort_results=True,
                 plot_prcurve=False,
                 plot_dir='',
                 suffix='_0',
                 gt_names=('boxes', 'classes'),
                 pred_names=('pred_boxes', 'pred_classes', 'pred_scores'),
                 det_eval_cls=-1):
        self.label_names = label_names
        self.iou_threshs = iou_threshs
        self.ignore_thresh = ignore_thresh
        self.ignore_mode = ignore_mode
        self.sort_results = sort_results
        self.gt_names = gt_names
        self.pred_names = pred_names
        self.plot_prcurve = plot_prcurve
        self.plot_dir = plot_dir
        self.suffix = suffix
        self.labels = None
        self.det_eval_cls=det_eval_cls


    def process_results(self, results):
        if self.sort_results:
            for i in range(len(results)):
                result = results[i]
                pred_boxes = result[self.pred_names[0]]
                pred_classes = result[self.pred_names[1]]
                pred_scores = result[self.pred_names[2]]

                if self.det_eval_cls != -1 and len(pred_classes) > 0:
                    keep = np.where(pred_classes == self.det_eval_cls)[0]
                    pred_boxes = pred_boxes[keep, ...]
                    pred_classes = pred_classes[keep, ...]
                    pred_scores  = pred_scores[keep, ...]
                sorted_ind = np.argsort(-pred_scores)
                result[self.pred_names[0]] = pred_boxes[sorted_ind]
                result[self.pred_names[1]] = pred_classes[sorted_ind]
                result[self.pred_names[2]] = pred_scores[sorted_ind]
        return results


    def eval_all_classes(self, labels, results):
        assert len(labels) == len(results)
        all_rets = []
        num_images = len(results)

        for c in range(1, len(self.label_names) + 1):
            gt_boxes_list = []
            pred_boxes_list = []
            pred_scores_list = []
            ignore_boxes_list = []
            for i in range(num_images):
                label = labels[i]
                boxes = label[self.gt_names[0]].astype(np.float32)
                classes = label[self.gt_names[1]].astype(np.int32)
                gt_boxes_list.append(boxes[classes == c])
                ignore_boxes_list.append(boxes[classes == -c])
                result = results[i]
                pred_boxes = result[self.pred_names[0]].astype(np.float32)
                pred_classes = result[self.pred_names[1]].astype(np.int32)
                pred_scores = result[self.pred_names[2]].astype(np.float32)
                pred_boxes_list.append(pred_boxes[pred_classes == c])
                pred_scores_list.append(pred_scores[pred_classes == c])
            all_rets.append(self.eval_single_class(gt_boxes_list,
                                                   pred_boxes_list,
                                                   pred_scores_list,
                                                   ignore_boxes_list,
                                                   eval_cls=c))
        recall = [{} for _ in self.iou_threshs]
        precision = [{} for _ in self.iou_threshs]
        ap = [{} for _ in self.iou_threshs]
        for i, label_name in enumerate(self.label_names):
            for iou_idx in range(len(self.iou_threshs)):
                recall[iou_idx][label_name], precision[iou_idx][label_name], \
                    ap[iou_idx][label_name] = all_rets[i][iou_idx]
        return recall, precision, ap


    def eval_single_class(self, gt_boxes_list, pred_boxes_list, pred_scores_list, ignore_boxes_list=None,eval_cls=0):
        assert len(gt_boxes_list) == len(pred_boxes_list) == len(pred_scores_list)
        tps_list = []
        fps_list = []
        npos = 0
        num_images = len(gt_boxes_list)
        for i in range(num_images):
            gt_boxes = gt_boxes_list[i]
            pred_boxes = pred_boxes_list[i]
            ignore_boxes = ignore_boxes_list[i] if ignore_boxes_list is not None else None
            tps, fps = self.eval_single_class_single_image(gt_boxes, pred_boxes, ignore_boxes)
            tps_list.append(tps)
            fps_list.append(fps)
            npos += len(gt_boxes)
        tps = np.concatenate(tps_list, axis=1)
        fps = np.concatenate(fps_list, axis=1)
        pred_scores = np.concatenate(pred_scores_list, axis=0)
        sorted_ind = np.argsort(-pred_scores)
        tps = tps[:, sorted_ind]
        fps = fps[:, sorted_ind]
        ret = []
        for iou_idx in range(len(self.iou_threshs)):
            tp = np.cumsum(tps[iou_idx])
            fp = np.cumsum(fps[iou_idx])
            recall = tp / max(float(npos), 1)
            precision = tp / np.maximum(tp + fp, np.finfo(np.float32).eps)
            if self.plot_prcurve and len(recall) > 0 and len(precision) > 0:
                cur_pid = self.suffix
                cls_id = '_clsid-' + str(eval_cls)
                iou_id = '_iou-' + str(self.iou_threshs[iou_idx])
                save_name = "prcurve{}{}{}.png".format(cur_pid,cls_id,iou_id)
                self.plot_score(os.path.join(self.plot_dir, save_name),recall.tolist(), {'precision':precision.tolist()}, x_name='recall', y_name='precision')
            ap = utils.average_precision(recall, precision)
            recall = recall[-1] if len(recall) > 0 else 0
            precision = precision[-1] if len(precision) > 0 else 0
            ret.append((recall, precision, ap))
        return ret


    def eval_single_class_single_image(self, gt_boxes, pred_boxes, ignore_boxes=None):
        tps = np.zeros((len(self.iou_threshs), len(pred_boxes)), dtype=np.float32)
        fps = np.zeros((len(self.iou_threshs), len(pred_boxes)), dtype=np.float32)
        if len(pred_boxes) == 0:
            return tps, fps
        if len(gt_boxes) > 0:
            gt_overlaps = utils.overlap(pred_boxes, gt_boxes)
        if ignore_boxes is not None and len(ignore_boxes) > 0:
            ignore_overlaps = utils.overlap(pred_boxes, ignore_boxes, mode=self.ignore_mode)
        match_gts = np.zeros((len(self.iou_threshs), len(gt_boxes)), dtype=np.bool_)
        for i in range(len(pred_boxes)):
            iou_max = -np.inf
            j_max = -1
            for j in range(len(gt_boxes)):
                iou = gt_overlaps[i, j]
                if iou > iou_max:
                    iou_max = iou
                    j_max = j
            match_ignore = False
            if ignore_boxes is not None and len(ignore_boxes) > 0:
                for k in range(len(ignore_boxes)):
                    ignore_iou = ignore_overlaps[i, k]
                    if ignore_iou > self.ignore_thresh:
                        match_ignore = True
                        break
            for iou_idx, iou_thresh in enumerate(self.iou_threshs):
                # if match a gt, tp = 1, fp = 0
                # if match an ignored gt, tp = 0, fp = 0
                # if match no gt, tp = 0, fp = 1
                if iou_max > iou_thresh and not match_gts[iou_idx, j_max]:
                    tps[iou_idx, i] = 1.0
                    match_gts[iou_idx, j_max] = True
                elif not match_ignore:
                    fps[iou_idx, i] = 1.0
        return tps, fps



class WiderEvaluator(object):
    def __init__(self, anno_path,
                 iou_thresh_list=[0.5],
                 det_eval_cls=-1,
                 kps_eval_cls=-1,
                 det_pred_names=('pred_boxes', 'pred_classes', 'pred_scores'),
                 kps_pred_names=('pred_keypoints', 'pred_kps_scores')):
        self.anno_path = anno_path
        self.iou_thresh_list = iou_thresh_list
        self.det_eval_cls = det_eval_cls
        self.kps_eval_cls = kps_eval_cls
        self.det_pred_names = det_pred_names
        self.kps_pred_names = kps_pred_names


    def eval_det(self, results):
        """evaluate widerface mAP"""

        pred_res = copy.deepcopy(results)
        pred_res = utils.parse_pred(pred_res, self.det_pred_names, self.det_eval_cls)
        utils.norm_score(pred_res)
        facebox_list, event_list, file_list, hard_gt_list, medium_gt_list, easy_gt_list = utils.get_gt_boxes(self.anno_path)
        event_num = len(event_list)
        thresh_num = 1000
        settings = ['easy', 'medium', 'hard']
        setting_gts = [easy_gt_list, medium_gt_list, hard_gt_list]
        aps_list = []
        for iou_idx, iou_thresh in enumerate(self.iou_thresh_list):
            print("widerface test: %s/%s-th iou_thresh:%s" % (iou_idx, len(self.iou_thresh_list), iou_thresh))
            aps = []
            for setting_id in range(3):
                # different setting
                gt_list = setting_gts[setting_id]
                count_face = 0
                pr_curve = np.zeros((thresh_num, 2)).astype('float')
                # [hard, medium, easy]
                # pbar = tqdm.tqdm()
                for i in range(event_num):
                    # pbar.set_description('Processing %s' % (settings[setting_id]))
                    event_name = str(event_list[i][0][0])
                    img_list = file_list[i][0]
                    pred_list = pred_res[event_name]
                    sub_gt_list = gt_list[i][0]
                    gt_bbx_list = facebox_list[i][0]

                    for j in range(len(img_list)):
                        pred_info = pred_list[str(img_list[j][0][0])]

                        gt_boxes = gt_bbx_list[j][0].astype('float')
                        keep_index = sub_gt_list[j][0]
                        count_face += len(keep_index)

                        if len(gt_boxes) == 0 or len(pred_info) == 0:
                            continue
                        ignore = np.zeros(gt_boxes.shape[0])
                        if len(keep_index) != 0:
                            ignore[keep_index - 1] = 1
                        pred_recall, proposal_list = utils.image_eval(pred_info, gt_boxes, ignore, iou_thresh)

                        _img_pr_info = utils.img_pr_info(thresh_num, pred_info, proposal_list, pred_recall)

                        pr_curve += _img_pr_info

                pr_curve = utils.dataset_pr_info(thresh_num, pr_curve, count_face)

                propose = pr_curve[:, 0]
                recall = pr_curve[:, 1]

                ap = utils.voc_ap(recall, propose)
                aps.append(ap)

            print("==================== Results ====================")
            print("Easy   Val AP: %.5f" % (aps[0]))
            print("Medium Val AP: %.5f" % (aps[1]))
            print("Hard   Val AP: %.5f" % (aps[2]))
            print("=================================================")



if __name__ == '__main__':

    flag_select = 'person' # face_0520 / widerface /person / pet

    if flag_select == 'face_0520':
        # face_0520 test
        pred_pkl_path = 'face0520/pred_100_results.pkl'
        pred_result = pickle.load(open(pred_pkl_path, 'rb'))
        anno_pkl_path = 'face0520/anno_100_results.pkl'
        labels = pickle.load(open(anno_pkl_path, 'rb'))

        map_Evaluator = MAPEvaluator(iou_threshs= [0.5], label_names=['face']) 
        results = map_Evaluator.process_results(pred_result)
        recall, _, ap = map_Evaluator.eval_all_classes(labels, results)
        
        print('face_0520: ap /ar:', ap[0]['face'], '/', recall[0]['face'])
    elif flag_select == 'widerface':
        # widerface test
        pred_pkl_path = 'widerface/widerface_pred_results.pkl'
        pred_result = pickle.load(open(pred_pkl_path, 'rb'))

        gt_dir ='widerface/ground_truth/'
        face_Evaluator = WiderEvaluator(gt_dir, iou_thresh_list=[0.5])
        face_Evaluator.eval_det(pred_result)
    elif flag_select == 'person':
        pred_pkl_path = 'person/person_pred_100_results.pkl'
        pred_result = pickle.load(open(pred_pkl_path, 'rb'))
        anno_pkl_path = 'person/person_anno_100_results.pkl'
        labels = pickle.load(open(anno_pkl_path, 'rb'))

        map_Evaluator = MAPEvaluator(iou_threshs= [0.5], label_names=['_', '_', 'person', '_', '_', '_', '_', 'pet']) 
        results = map_Evaluator.process_results(pred_result)
        recall, _, ap = map_Evaluator.eval_all_classes(labels, results)
        
        print('person: ap /ar:', ap[0]['person'], '/', recall[0]['person'])
    elif flag_select == 'pet':
        pred_pkl_path = 'pet/pet_pred_results.pkl'
        pred_result = pickle.load(open(pred_pkl_path, 'rb'))
        anno_pkl_path = 'pet/pet_anno_results.pkl'
        labels = pickle.load(open(anno_pkl_path, 'rb'))

        map_Evaluator = MAPEvaluator(iou_threshs= [0.5], label_names=['_', '_', 'person', '_', '_', '_', '_', 'pet']) 
        results = map_Evaluator.process_results(pred_result)
        recall, _, ap = map_Evaluator.eval_all_classes(labels, results)
        
        print('pet: ap /ar:', ap[0]['pet'], '/', recall[0]['pet'])


