import numpy as np
from scipy.spatial import cKDTree
from utils.pts import PtsUtil

class ReconstructionUtil:
    
    @staticmethod
    def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01):
        kdtree = cKDTree(combined_point_cloud)
        distances, _ = kdtree.query(target_point_cloud)
        covered_points = np.sum(distances < threshold)
        coverage_rate = covered_points / target_point_cloud.shape[0]
        return coverage_rate
    
    @staticmethod
    def compute_overlap_rate(new_point_cloud, combined_point_cloud, threshold=0.01):
        kdtree = cKDTree(combined_point_cloud)
        distances, _ = kdtree.query(new_point_cloud)
        overlapping_points = np.sum(distances < threshold)
        overlap_rate = overlapping_points / new_point_cloud.shape[0]
        return overlap_rate
    
    @staticmethod
    def combine_point_with_view_sequence(point_list, view_sequence):
        selected_views = []
        for view_index, _ in view_sequence:
            selected_views.append(point_list[view_index])
        return np.vstack(selected_views)
    
    @staticmethod 
    def compute_next_view_coverage_list(views, combined_point_cloud, target_point_cloud, threshold=0.01):
        best_view = None
        best_coverage_increase = -1
        current_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold)
        
        for view_index, view in enumerate(views):
            candidate_views = combined_point_cloud + [view]
            down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(candidate_views, threshold)
            new_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, down_sampled_combined_point_cloud, threshold)
            coverage_increase = new_coverage - current_coverage
            if coverage_increase > best_coverage_increase:
                best_coverage_increase = coverage_increase
                best_view = view_index
        return best_view, best_coverage_increase

    
    @staticmethod
    def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, threshold=0.01, overlap_threshold=0.3, status_info=None):
        selected_views = []
        current_coverage = 0.0
        remaining_views = list(range(len(point_cloud_list)))
        view_sequence = []
        cnt_processed_view = 0
        while remaining_views:
            best_view = None
            best_coverage_increase = -1
        
            for view_index in remaining_views:
                
                if selected_views:
                    combined_old_point_cloud = np.vstack(selected_views)
                    down_sampled_old_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_old_point_cloud,threshold)
                    down_sampled_new_view_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud_list[view_index],threshold)
                    overlap_rate = ReconstructionUtil.compute_overlap_rate(down_sampled_new_view_point_cloud,down_sampled_old_point_cloud, threshold)
                    if overlap_rate < overlap_threshold:
                        continue
                    
                candidate_views = selected_views + [point_cloud_list[view_index]]
                combined_point_cloud = np.vstack(candidate_views)
                down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold)
                new_coverage = ReconstructionUtil.compute_coverage_rate(target_point_cloud, down_sampled_combined_point_cloud, threshold)
                coverage_increase = new_coverage - current_coverage
                #print(f"view_index: {view_index}, coverage_increase: {coverage_increase}")
                if coverage_increase > best_coverage_increase:
                    best_coverage_increase = coverage_increase
                    best_view = view_index
                    cnt_processed_view += 1
                    if status_info is not None:
                        
                        sm = status_info["status_manager"]
                        app_name = status_info["app_name"]
                        runner_name = status_info["runner_name"]
                        sm.set_status(app_name, runner_name, "current coverage", current_coverage)
                        sm.set_progress(app_name, runner_name, "processed view", cnt_processed_view, len(point_cloud_list))
            
            if best_view is not None:
                if best_coverage_increase <=1e-3:
                    break
                selected_views.append(point_cloud_list[best_view])
                remaining_views.remove(best_view)
                if best_coverage_increase > 0:
                    current_coverage += best_coverage_increase
                
                view_sequence.append((best_view, current_coverage))
                
            else:
                break
        if status_info is not None:
            sm = status_info["status_manager"]
            app_name = status_info["app_name"]
            runner_name = status_info["runner_name"]
            sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list))
        return view_sequence, remaining_views
    
    @staticmethod
    def filter_points(points, points_normals, cam_pose,  voxel_size=0.005, theta=45):
        sampled_points = PtsUtil.voxel_downsample_point_cloud(points, voxel_size)
        kdtree = cKDTree(points_normals[:,:3])
        _, indices = kdtree.query(sampled_points)
        nearest_points = points_normals[indices]

        normals = nearest_points[:, 3:]
        camera_axis = -cam_pose[:3, 2] 
        normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
        cos_theta = np.dot(normals_normalized, camera_axis)
        theta_rad = np.deg2rad(theta)
        filtered_sampled_points= sampled_points[cos_theta > np.cos(theta_rad)]
        
        return filtered_sampled_points[:, :3]