# Copyright (c) OpenMMLab. All rights reserved.
# Adapt from
# https://github.com/open-mmlab/mmpose/blob/master/mmpose/datasets/datasets/base/kpt_2d_sview_rgb_img_top_down_dataset.py
import numpy as np

from easycv.framework.errors import KeyError
from .base_evaluator import Evaluator
from .builder import EVALUATORS
from .metric_registry import METRICS
from .top_down_eval import (keypoint_auc, keypoint_epe, keypoint_nme,
                            keypoint_pck_accuracy)


@EVALUATORS.register_module
class KeyPointEvaluator(Evaluator):
    """ KeyPoint evaluator.
    """

    def __init__(self,
                 dataset_name=None,
                 metric_names=['PCK', 'PCKh', 'AUC', 'EPE', 'NME'],
                 pck_thr=0.2,
                 pckh_thr=0.7,
                 auc_nor=30):
        """

        Args:
            dataset_name: eval dataset name
            metric_names: eval metrics name
            pck_thr (float): PCK threshold, default as 0.2.
            pckh_thr (float): PCKh threshold, default as 0.7.
            auc_nor (float): AUC normalization factor, default as 30 pixel.
        """
        super(KeyPointEvaluator, self).__init__(dataset_name, metric_names)
        self._pck_thr = pck_thr
        self._pckh_thr = pckh_thr
        self._auc_nor = auc_nor
        self.dataset_name = dataset_name
        allowed_metrics = ['PCK', 'PCKh', 'AUC', 'EPE', 'NME']
        for metric in metric_names:
            if metric not in allowed_metrics:
                raise KeyError(f'metric {metric} is not supported')

    def _evaluate_impl(self, preds, coco_db, **kwargs):
        ''' keypoint evaluation code which will be run after
        all test batched data are predicted

        Args:
            preds: dict with key ``keypoints`` whose shape is Nx3
            coco_db: the db of wholebody coco datasource, sorted by 'bbox_id'

        Return:
            a dict,  each key is metric_name, value is metric value
        '''
        assert len(preds) == len(coco_db)
        eval_res = {}

        outputs = []
        gts = []
        masks = []
        box_sizes = []
        threshold_bbox = []
        threshold_head_box = []

        for pred, item in zip(preds, coco_db):
            outputs.append(np.array(pred['keypoints'])[:, :-1])
            gts.append(np.array(item['joints_3d'])[:, :-1])
            masks.append((np.array(item['joints_3d_visible'])[:, 0]) > 0)
            if 'PCK' in self.metric_names:
                bbox = np.array(item['bbox'])
                bbox_thr = np.max(bbox[2:])
                threshold_bbox.append(np.array([bbox_thr, bbox_thr]))
            if 'PCKh' in self.metric_names:
                head_box_thr = item['head_size']
                threshold_head_box.append(
                    np.array([head_box_thr, head_box_thr]))
            box_sizes.append(item.get('box_size', 1))

        outputs = np.array(outputs)
        gts = np.array(gts)
        masks = np.array(masks)
        threshold_bbox = np.array(threshold_bbox)
        threshold_head_box = np.array(threshold_head_box)
        box_sizes = np.array(box_sizes).reshape([-1, 1])

        if 'PCK' in self.metric_names:
            _, pck, _ = keypoint_pck_accuracy(outputs, gts, masks,
                                              self._pck_thr, threshold_bbox)
            eval_res['PCK'] = pck

        if 'PCKh' in self.metric_names:
            _, pckh, _ = keypoint_pck_accuracy(outputs, gts, masks,
                                               self._pckh_thr,
                                               threshold_head_box)
            eval_res['PCKh'] = pckh

        if 'AUC' in self.metric_names:
            eval_res['AUC'] = keypoint_auc(outputs, gts, masks, self._auc_nor)

        if 'EPE' in self.metric_names:
            eval_res['EPE'] = keypoint_epe(outputs, gts, masks)

        if 'NME' in self.metric_names:
            normalize_factor = self._get_normalize_factor(
                gts=gts, box_sizes=box_sizes)
            eval_res['NME'] = keypoint_nme(outputs, gts, masks,
                                           normalize_factor)
        return eval_res

    def _get_normalize_factor(self, gts, *args, **kwargs):
        """Get the normalize factor. generally inter-ocular distance measured
        as the Euclidean distance between the outer corners of the eyes is
        used. This function should be overrode, to measure NME.

        Args:
            gts (np.ndarray[N, K, 2]): Groundtruth keypoint location.

        Returns:
            np.ndarray[N, 2]: normalized factor
        """
        return np.ones([gts.shape[0], 2], dtype=np.float32)


METRICS.register_default_best_metric(KeyPointEvaluator, 'PCK', 'max')