import torch
import annotations.stereotype as stereotype
from utils.pose_util import PoseUtil


@stereotype.evaluation_method("delta_pose")
def evaluate(output_list, data_list):
    results = {"scalars": {}}
    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 = rotation_angle_distance(gt_delta_rot_mat, est_delta_rot_mat)
        rot_angle_list.extend(list(rotation_angles))

    results["scalars"]["delta_rotation"] = float(sum(rot_angle_list) / len(rot_angle_list))
    return results


def rotation_angle_distance(R1, R2):
    R = torch.matmul(R1, R2.transpose(1, 2))
    trace = torch.diagonal(R, dim1=1, dim2=2).sum(-1)
    angle = torch.acos(torch.clamp((trace - 1) / 2, -1.0, 1.0))/torch.pi*180
    return angle