from utils.pose import PoseUtil
import PytorchBoot.stereotype as stereotype
import PytorchBoot.namespace as namespace

@stereotype.evaluation_method("delta_pose_diff")
class DeltaPoseDiff:
    def __init__(self, _):
        pass
    def evaluate(self, output_list, data_list):
        results = {namespace.TensorBoard.SCALAR: {}}
        rot_angle_list = []
        for output, data in zip(output_list, data_list):
            gt_delta_rot_6d = data['delta_rot_6d']
            est_delta_rot_6d = output['estimated_delta_rot_6d']
            gt_delta_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(gt_delta_rot_6d)
            est_delta_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(est_delta_rot_6d)
            rotation_angles = PoseUtil.rotation_angle_distance(gt_delta_rot_mat, est_delta_rot_mat)
            rot_angle_list.extend(list(rotation_angles))

        results[namespace.TensorBoard.SCALAR]["delta_rotation"] = float(sum(rot_angle_list) / len(rot_angle_list))
        return results
    


@stereotype.evaluation_method("coverage_rate_increase",comment="unfinished")
class ConverageRateIncrease:
    def __init__(self, config):
        pass

    def evaluate(self, output_list, data_list):
        return