"""
Evaluation metrics of pose estimation models
"""
import tensorflow as tf


class InterhandMetric:
    """Interhand model evaluation metric
    """
    def __init__(self, handness_thresholds=[0.5], log_residuals=False):
        """
        Args:
            handness_thresholds (list, optional): handedness classification thresholds. Defaults to [0.5].
            log_residuals (bool, optional): flag to keep track of keypoint residuals. Defaults to False.
        """
        # handedness metric, relative depth evaluation metric, 2D keypoint metric, keypoint depth metric
        self.handness_metric = InterhandHandnessMetric(handness_thresholds)
        self.depth_metric = tf.keras.metrics.MeanSquaredError()
        self.joints_2d_metric = tf.keras.metrics.MeanAbsoluteError()
        self.joints_depth_metric = tf.keras.metrics.MeanSquaredError()
        # if logging -> setup logging lists
        self.log_residuals = log_residuals
        if self.log_residuals:
            self.joint_2d_residuals = list()
            self.joint_depth_residuals = list()
            self.depth_residuals = list()

    def update_state(self, y_true, y_pred):
        """Update metric state

        Args:
            y_true (dict): true label dictionary
            y_pred (list): list of model predictions - handedness, keypoint estimation, relative root depth
        """
        handness_pred, joints_pred, depth_pred = y_pred
        handness_true = y_true['handness']
        joints_true = y_true['joints']
        depth_true = y_true['relative_root_depth']
        # update handedness metric
        self.handness_metric.update_state(handness_true, handness_pred)
        # interacting hands mask
        two_hands_mask = tf.reduce_all(tf.cast(handness_true, tf.bool), axis=1)
        # eval depth only if hands are interacting
        depth_true = tf.boolean_mask(depth_true, two_hands_mask)
        depth_pred = tf.boolean_mask(depth_pred, two_hands_mask)
        if tf.size(depth_pred):
            # update relative root depth metric 
            self.depth_metric.update_state(depth_true, depth_pred)
            # if logging -> calculate and log residuals
            if self.log_residuals:
                residuals = tf.reshape(depth_true - depth_pred, (-1, )).numpy().tolist()
                self.depth_residuals += residuals
        # masked interacting keypoints
        two_hands_joints_true = tf.boolean_mask(joints_true, two_hands_mask, axis=0)
        two_hands_joints_pred = tf.boolean_mask(joints_pred, two_hands_mask, axis=0)
        # masked single hand keypoints
        single_hand_joints_true = tf.boolean_mask(joints_true, tf.logical_not(two_hands_mask), axis=0)
        single_hand_joints_pred = tf.boolean_mask(joints_pred, tf.logical_not(two_hands_mask), axis=0)

        single_hand_handness = tf.boolean_mask(handness_true, tf.logical_not(two_hands_mask), axis=0)
        # single hand keypoint estimation metric
        if tf.size(single_hand_handness) > 0:
            # mask unused hand keypoints
            single_hand_mask = tf.cast(tf.repeat(single_hand_handness, 21, axis=1), tf.bool)
            single_hand_joints_true = tf.boolean_mask(single_hand_joints_true, single_hand_mask)
            single_hand_joints_pred = tf.boolean_mask(single_hand_joints_pred, single_hand_mask)
            # update keypoint depth metric and 2D metric
            self.joints_depth_metric.update_state(single_hand_joints_true[..., 2], single_hand_joints_pred[..., 2])
            self.joints_2d_metric.update_state(single_hand_joints_true[..., :2], single_hand_joints_pred[..., :2])
            # if logging -> calculate and log residuals
            if self.log_residuals:
                residuals = tf.reshape(single_hand_joints_true[..., 2] - single_hand_joints_pred[..., 2], (-1, )).numpy().tolist()
                self.joint_depth_residuals += residuals
                residuals = tf.reshape(single_hand_joints_true[..., :2] - single_hand_joints_pred[..., :2], (-1, )).numpy().tolist()
                self.joint_2d_residuals += residuals

        # interacting hands keypoint estimation metric
        if tf.size(two_hands_joints_pred) > 0:
            # update keypoint depth metric and 2D metric
            self.joints_2d_metric.update_state(two_hands_joints_true[..., :2], two_hands_joints_pred[..., :2])
            self.joints_depth_metric.update_state(two_hands_joints_true[..., 2], two_hands_joints_pred[..., 2])
            # if logging -> calculate and log residuals
            if self.log_residuals:
                residuals = tf.reshape(two_hands_joints_true[..., 2] - two_hands_joints_pred[..., 2], (-1, )).numpy().tolist()
                self.joint_depth_residuals += residuals
                residuals = tf.reshape(two_hands_joints_true[..., :2] - two_hands_joints_pred[..., :2], (-1, )).numpy().tolist()
                self.joint_2d_residuals += residuals

    def reset_state(self):
        """Reset metric state
        """
        self.handness_metric.reset_state()
        self.depth_metric.reset_state()
        self.joints_2d_metric.reset_state()
        self.joints_depth_metric.reset_state()

        if self.log_residuals:
            self.joint_2d_residuals = list()
            self.joint_depth_residuals = list()
            self.depth_residuals = list()

    def result(self):
        """Calculate metric results

        Returns:
            float, float, float, float: results of individual metrics: handedness, relative root depth, 2D keypoints, depth keypoints
        """
        handness_result = self.handness_metric.result()
        depth_result = self.depth_metric.result()
        joints_2d_metric = self.joints_2d_metric.result()
        joints_depth_metric = self.joints_depth_metric.result()
        return handness_result, depth_result, joints_2d_metric, joints_depth_metric

    def get_residuals(self):
        """Get logged residuals

        Returns:
            dict: dictionary of logged residuals
        """
        # if not logging return None
        if not self.log_residuals:
            return None
        return {
            'joints_2d': self.joint_2d_residuals,
            'joints_depth': self.joint_depth_residuals,
            'depth': self.depth_residuals
            }


class InterhandHandnessMetric:
    """Interhand handedness metric
    """
    def __init__(self, thresholds):
        """
        Args:
            thresholds (str): classification thresholds
        """
        self.thresholds = tf.constant(thresholds, dtype=tf.float32)
        # (tholds, true, predicted)
        self.confusion_matrix = tf.zeros((tf.shape(self.thresholds)[0], 3, 3), dtype=tf.int32)
        # (tholds, class, [TP, FP, TN, FN])
        self.stats = tf.zeros((tf.shape(self.thresholds)[0], 3, 4), dtype=tf.float32)

    def update_state(self, y_true, y_pred):
        """Update metric state

        Args:
            y_true (tf.Tensor): true handedness
            y_pred (tf.Tensor): predicted handedness
        """
        # for each threshold
        for i in tf.range(tf.shape(self.thresholds)[0]):
            # sparse true and predicted classes
            class_true = self.dense_label_to_sparse(y_true)
            class_pred = self.dense_label_to_sparse(y_pred_discrete)
            # thresholded prediction
            y_pred_discrete = tf.cast(y_pred > self.thresholds[i], tf.float32)
            # three classes
            for cls in tf.range(3):
                # prepare stats masks
                true_mask = class_true == cls
                false_mask = class_true != cls
                positives_mask = class_pred == cls
                negatives_mask = class_pred != cls
                # calculate TP, FP, TN and FN
                true_positives = tf.reduce_sum(tf.cast(tf.logical_and(true_mask, positives_mask), tf.float32))
                false_positives = tf.reduce_sum(tf.cast(tf.logical_and(false_mask, positives_mask), tf.float32))
                true_negatives = tf.reduce_sum(tf.cast(tf.logical_and(true_mask, negatives_mask), tf.float32))
                false_negatives = tf.reduce_sum(tf.cast(tf.logical_and(false_mask, negatives_mask), tf.float32))
                # update stats
                updated_stats = self.stats[i, cls, :] + [true_positives, false_positives, true_negatives, false_negatives]
                self.stats = tf.tensor_scatter_nd_update(self.stats, [[i, cls, 0], [i, cls, 1], [i, cls, 2], [i, cls, 3]], updated_stats)
                # update confusion matrix
                cls_pred = tf.cast(tf.boolean_mask(class_pred, true_mask), tf.int32)
                updated_cm = tf.math.bincount(cls_pred, minlength=3) + self.confusion_matrix[i, cls, :]
                self.confusion_matrix = tf.tensor_scatter_nd_update(self.confusion_matrix, [[i, cls, 0], [i, cls, 1], [i, cls, 2]], updated_cm)


    def reset_state(self):
        """Reset metric state
        """
        self.confusion_matrix = tf.zeros((3, 3), dtype=tf.float32)
        self.stats = tf.zeros((4, 3), dtype=tf.float32)

    def result(self):
        """Return metric results

        Returns:
            np.ndarray, np.ndarray: TP, FP, TN and FN stats, confusion matrix. Results are calculated for every threshold.
        """
        return self.stats.numpy(), self.confusion_matrix.numpy()

    @staticmethod
    def dense_label_to_sparse(dense_label):
        """Convert dense multiclass vector to a sparse class

        Args:
            dense_label (tf.Tensor): N dense labels [N, 2]

        Returns:
            tf.Tensor: N sparse labels [N, ]
        """
        # input labels are two indicators [left, right], convert them to a sparse labels 0, 1, 2 (3 combinations - at least one is expected to be positive)
        # add weights to indicators and sum them
        normed = dense_label * tf.constant([2, 1], dtype=tf.float32)
        sparse = tf.reduce_sum(normed, axis=1) - 1
        return tf.cast(sparse, tf.int32)


class FaceKeypointMetric:
    """Face keypoint estimation metric
    """
    def __init__(self):
        # examples are tagged with conditions
        # evaluate examples overall and individually based on tags
        # same metric for all possible tags
        metric = tf.keras.metrics.MeanAbsoluteError

        self.overall_metric = metric(name='overall')
        self.attrib_metrics = [
            metric(name='Pose'),
            metric(name='Expression'),
            metric(name='Illumination'),
            metric(name='Make-up'),
            metric(name='Occlusion'),
            metric(name='Blur'),
        ]
    
    def update_state(self, y_true, y_pred):
        """Update metric state

        Args:
            y_true (dict): true labels dictionary
            y_pred (tf.Tensor): predicted keypoints
        """
        keypoints_true = y_true['keypoints']
        attributes = y_true['attribute']
        # update overall metric
        keypoints_pred = y_pred
        self.overall_metric.update_state(keypoints_true, keypoints_pred)
        # update other metrics (based on its tags)
        for i, metric in enumerate(self.attrib_metrics):
            # mask predictions and labels based on tag and update metric
            mask = tf.cast(attributes[:, i], tf.bool)
            masked_true = tf.boolean_mask(keypoints_true, mask, axis=0)
            masked_pred = tf.boolean_mask(keypoints_pred, mask, axis=0)
            metric.update_state(masked_true, masked_pred)

    def reset_state(self):
        """Reset metric state
        """
        self.overall_metric.reset_state()
        for metric in self.attrib_metrics:
            metric.reset_state()

    def result(self):
        """Calculate metric results

        Returns:
            float, dict: overall result and results of tagged metrics
        """
        overall_result = self.overall_metric.result()
        attrib_results = {metric.name: metric.result() for metric in self.attrib_metrics}
        return overall_result, attrib_results


class HandKeypointMetric:
    """Hand pose estimation evaluation metric
    """
    def __init__(self, thresholds=[0.5]):
        """
        Args:
            thresholds (list, optional): list of classification head thresholds. Defaults to [0.5].
        """
        # custom hand classification metric, AuROC classification metric, MAE for single and interacting poses
        self.handness_metric = HandClassificationMetric(thresholds)
        self.auc = tf.keras.metrics.AUC(num_thresholds=100)
        self.single_hand_metric = tf.keras.metrics.MeanAbsoluteError(name='single_hand')
        self.interacting_metric = tf.keras.metrics.MeanAbsoluteError(name='interacting')

    def update_state(self, y_true, y_pred):
        handness_true = y_true['class']
        handness_pred = y_pred[0]
        # update classification metrics
        self.handness_metric.update_state(handness_true, handness_pred)
        self.auc.update_state(handness_true, handness_pred)

        keypoints_true = y_true['keypoints']
        single_hand_keypoints_pred = y_pred[1]
        interacting_keypoints_pred = y_pred[2]
        # mask single hand examples
        single_hand_mask = tf.reshape(tf.logical_not(tf.cast(handness_true, tf.bool)), (-1, ))
        single_hand_keypoints_true = tf.boolean_mask(keypoints_true, single_hand_mask, axis=0)
        single_hand_keypoints_pred = tf.boolean_mask(single_hand_keypoints_pred, single_hand_mask, axis=0)
        # single hand metric update
        self.single_hand_metric.update_state(single_hand_keypoints_true[:, :21, :], single_hand_keypoints_pred)
        # mask interacting hand examples
        interacting_mask = tf.reshape(tf.cast(handness_true, tf.bool), (-1, ))
        interacting_keypoints_true = tf.boolean_mask(keypoints_true, interacting_mask, axis=0)
        interacting_keypoints_pred = tf.boolean_mask(interacting_keypoints_pred, interacting_mask, axis=0)
        # interacting hands metric update
        self.interacting_metric.update_state(interacting_keypoints_true, interacting_keypoints_pred)

    def reset_state(self):
        """Reset metric state
        """
        self.handness_metric.reset_state()
        self.auc.reset_state()
        self.single_hand_metric.reset_state()
        self.interacting_metric.reset_state()

    def result(self):
        """Calculate results

        Returns:
            dict: dictionary of evaluation results
        """
        handness = self.handness_metric.result()
        auc = self.auc.result()
        single_hand = self.single_hand_metric.result()
        interacting = self.interacting_metric.result()
        return {'handness': handness, 'area_under_roc': auc, 'single_hand': single_hand, 'interacting': interacting}


class HandClassificationMetric:
    """Custom hand classification metric
    """
    def __init__(self, thresholds):
        """
        Args:
            thresholds (list): list of classification head thresholds
        """
        self.thresholds = thresholds
        # precision and recall metric for every threshold
        self.handness_precisions = [tf.keras.metrics.Precision() for _ in thresholds]
        self.handness_recalls = [tf.keras.metrics.Recall() for _ in thresholds]

    def update_state(self, y_true, y_pred):
        """Update metric

        Args:
            y_true (tf.Tensor): true class
            y_pred (tf.Tensor): predicted class
        """
        # for every threshold update precision and recall
        for threshold, precision, recall in zip(self.thresholds, self.handness_precisions, self.handness_recalls):
            handness_pred = tf.cast(y_pred > threshold, tf.int32)
            precision.update_state(y_true, handness_pred)
            recall.update_state(y_true, handness_pred)

    def reset_state(self):
        """Reset metric state
        """
        for precision, recall in zip(self.handness_precisions, self.handness_recalls):
            precision.reset_state()
            recall.reset_state()

    def result(self):
        """Calcualte metric results

        Returns:
            dict: dictionary of metric results
        """
        results = list()
        # for each threshold calculate F1 score and output it with precision and recall
        for threshold, precision, recall in zip(self.thresholds, self.handness_precisions, self.handness_recalls):
            precision_result = precision.result()
            recall_result = recall.result()
            f1 = 2 * (precision_result * recall_result) / (precision_result + recall_result)
            results.append({
                'threshold': threshold,
                'precision': precision_result,
                'recall': recall_result,
                'f1': f1
            })
        return results
