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

class ReconstructionUtil:
    @staticmethod
    def reconstruct_with_pts(pts):
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(pts)
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
        mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9)
        densities = np.asarray(densities)
        vertices_to_remove = densities < np.quantile(densities, 0.2)
        mesh.remove_vertices_by_mask(vertices_to_remove)
        return mesh
    
    @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)
        #sampled_points = points
        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]
    
    @staticmethod
    def get_new_added_points(old_combined_pts, new_pts, threshold=0.005):
        if old_combined_pts.size == 0:
            return new_pts
        if new_pts.size == 0:
            return np.array([])

        tree = cKDTree(old_combined_pts)
        distances, _ = tree.query(new_pts, k=1)
        new_added_points = new_pts[distances > threshold]
        return new_added_points
    
if __name__ == "__main__":
    import os
    root = "/media/hofee/data/project/python/nbv_reconstruction/nbv_rec_visualize/data/sample/"
    name = "google_scan-box_0106"
    model_path = os.path.join(root, name, "sampled_model_points.txt")
    rec_path = os.path.join(root, name, "best_reconstructed_pts.txt")
    model_pts = np.loadtxt(model_path)   
    rec_pts = np.loadtxt(rec_path)
    import time
    start = time.time()
    model_mesh = ReconstructionUtil.reconstruct_with_pts(model_pts)
    end = time.time()
    print(f"Time taken to reconstruct model: {end-start}")
    rec_mesh = ReconstructionUtil.reconstruct_with_pts(rec_pts)
    output_dir = "/media/hofee/data/project/python/nbv_reconstruction/nbv_rec_visualize/mis/output_test"
    model_output_path = os.path.join(output_dir, f"{name}_model_mesh.obj")
    rec_output_path = os.path.join(output_dir, f"{name}_rec_mesh.obj")
    o3d.io.write_triangle_mesh(model_output_path, model_mesh)
    o3d.io.write_triangle_mesh(rec_output_path, rec_mesh)