import torch
import numpy as np
from utils.reconstruction import ReconstructionUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil
from utils.render import RenderUtil
import PytorchBoot.stereotype as stereotype
import PytorchBoot.namespace as namespace
from PytorchBoot.utils.log_util import Log



@stereotype.evaluation_method("pose_diff")
class PoseDiff:
    def __init__(self, _):
        pass
    
    def evaluate(self, output_list, data_list):
        results = {namespace.TensorBoard.SCALAR: {}}
        rot_angle_list = []
        trans_dist_list = []
        for output, data in zip(output_list, data_list):
            gt_pose_9d = data['best_to_world_pose_9d']
            pred_pose_9d = output['pred_pose_9d']
            gt_rot_6d = gt_pose_9d[:, :6]
            gt_trans = gt_pose_9d[:, 6:]
            pred_rot_6d = pred_pose_9d[:, :6]
            pred_trans = pred_pose_9d[:, 6:]
            gt_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(gt_rot_6d)
            pred_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_rot_6d)
            rotation_angles = PoseUtil.rotation_angle_distance(gt_rot_mat, pred_rot_mat)
            
            rot_angle_list.extend(list(rotation_angles))
            trans_dist = torch.norm(gt_trans-pred_trans, dim=1).mean().item()
            trans_dist_list.append(trans_dist)
            

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


@stereotype.evaluation_method("coverage_rate_increase")
class ConverageRateIncrease:
    def __init__(self, config):
        self.config = config
        self.renderer_path = config["renderer_path"]

    def evaluate(self, output_list, data_list):
        results = {namespace.TensorBoard.SCALAR: {}}
        gt_coverate_increase_list = []
        pred_coverate_increase_list = []
        cr_diff_list = []
        for output, data in zip(output_list, data_list):
            scanned_cr = data['scanned_coverage_rate']
            gt_cr = data["best_coverage_rate"]
            scene_path_list = data['scene_path']
            model_points_normals_list = data['model_points_normals']
            scanned_view_pts_list = data['scanned_target_pts_list']
            pred_pose_9ds = output['pred_pose_9d']
            nO_to_nL_pose_batch = data["nO_to_nL_pose"]
            voxel_threshold_list = data["voxel_threshold"]
            filter_degree_list = data["filter_degree"]
            first_frame_to_world = data["first_frame_to_world"]
            pred_n_to_world_pose_mats = torch.eye(4, device=pred_pose_9ds.device).unsqueeze(0).repeat(pred_pose_9ds.shape[0], 1, 1)
            pred_n_to_world_pose_mats[:,:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6])
            pred_n_to_world_pose_mats[:,:3,3] = pred_pose_9ds[:, 6:]
            pred_n_to_world_pose_mats = torch.matmul(first_frame_to_world, pred_n_to_world_pose_mats)
            for idx in range(len(scanned_cr)):
                model_points_normals = model_points_normals_list[idx]
                scanned_view_pts = scanned_view_pts_list[idx]
                voxel_threshold = voxel_threshold_list[idx]
                model_pts = model_points_normals[:,:3]
                down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
                old_scanned_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)
                gt_coverate_increase_list.append(gt_cr[idx]-old_scanned_cr)

                scene_path = scene_path_list[idx]
                pred_pose = pred_n_to_world_pose_mats[idx]
            
                filter_degree = filter_degree_list[idx]
                nO_to_nL_pose = nO_to_nL_pose_batch[idx]
                try:
                    new_pts, _ = RenderUtil.render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose)
                    pred_cr = self.compute_coverage_rate(scanned_view_pts, new_pts, down_sampled_model_pts, threshold=voxel_threshold)
                except Exception as e:
                    Log.warning(f"Error in scene {scene_path}, {e}")
                    pred_cr = old_scanned_cr
                pred_coverate_increase_list.append(pred_cr-old_scanned_cr)
                cr_diff_list.append(gt_cr[idx]-pred_cr)
                
        results[namespace.TensorBoard.SCALAR]["gt_cr_increase"] = float(sum(gt_coverate_increase_list) / len(gt_coverate_increase_list))
        results[namespace.TensorBoard.SCALAR]["pred_cr_increase"] = float(sum(pred_coverate_increase_list) / len(pred_coverate_increase_list))
        results[namespace.TensorBoard.SCALAR]["cr_diff"] = float(sum(cr_diff_list) / len(cr_diff_list))
        return results
    
    def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
        if new_pts is not None:
            new_scanned_view_pts = scanned_view_pts + [new_pts]
        else:
            new_scanned_view_pts = scanned_view_pts
        combined_point_cloud = np.vstack(new_scanned_view_pts)
        down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold)
        return ReconstructionUtil.compute_coverage_rate(model_pts, down_sampled_combined_point_cloud, threshold)