import numpy as np
import open3d as o3d
from scipy.spatial import cKDTree

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(point_cloud1, point_cloud2, threshold=0.01):
        kdtree1 = cKDTree(point_cloud1)
        kdtree2 = cKDTree(point_cloud2)
        distances1, _ = kdtree2.query(point_cloud1)
        distances2, _ = kdtree1.query(point_cloud2)
        overlapping_points1 = np.sum(distances1 < threshold)
        overlapping_points2 = np.sum(distances2 < threshold)
        
        overlap_rate1 = overlapping_points1 / point_cloud1.shape[0]
        overlap_rate2 = overlapping_points2 / point_cloud2.shape[0]
        
        return (overlap_rate1 + overlap_rate2) / 2
    
    @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 = ReconstructionUtil.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(target_point_cloud, point_cloud_list, threshold=0.01):
        selected_views = []
        current_coverage = 0.0
        remaining_views = list(range(len(point_cloud_list)))
        view_sequence = []
        target_point_cloud = ReconstructionUtil.downsample_point_cloud(target_point_cloud, threshold)
        while remaining_views:
            best_view = None
            best_coverage_increase = -1
            
            for view_index in remaining_views:
                candidate_views = selected_views + [point_cloud_list[view_index]]
                combined_point_cloud = np.vstack(candidate_views)
                
                down_sampled_combined_point_cloud = ReconstructionUtil.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
                
                if coverage_increase > best_coverage_increase:
                    best_coverage_increase = coverage_increase
                    best_view = view_index
            
            if best_view is not None:
                if best_coverage_increase <=1e-3: 
                    break
                selected_views.append(point_cloud_list[best_view])
                current_coverage += best_coverage_increase
                view_sequence.append((best_view, current_coverage))
                remaining_views.remove(best_view)
        return view_sequence, remaining_views
    
    
    @staticmethod
    def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, threshold=0.01, overlap_threshold=0.3):
        selected_views = []
        current_coverage = 0.0
        remaining_views = list(range(len(point_cloud_list)))
        view_sequence = []
        target_point_cloud = ReconstructionUtil.downsample_point_cloud(target_point_cloud, threshold)
        
        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 = ReconstructionUtil.downsample_point_cloud(combined_old_point_cloud,threshold)
                    down_sampled_new_view_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud_list[view_index],threshold)
                    overlap_rate = ReconstructionUtil.compute_overlap_rate(down_sampled_old_point_cloud,down_sampled_new_view_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 = ReconstructionUtil.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
            
            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
            
        return view_sequence, remaining_views
    
    
    def downsample_point_cloud(point_cloud, voxel_size=0.005):
        o3d_pc = o3d.geometry.PointCloud()
        o3d_pc.points = o3d.utility.Vector3dVector(point_cloud)
        downsampled_pc = o3d_pc.voxel_down_sample(voxel_size)
        return np.asarray(downsampled_pc.points)