import numpy as np
import open3d as o3d

def improved_multiway_registration(pcds, voxel_size=0.05, max_correspondence_distance_coarse=None, max_correspondence_distance_fine=None, overlap=3, quadratic_overlap=True, use_colored_icp=True):
    if max_correspondence_distance_coarse is None:
        max_correspondence_distance_coarse = voxel_size * 15
    if max_correspondence_distance_fine is None:
        max_correspondence_distance_fine = voxel_size * 1.5

    def preprocess_point_cloud(pcd, voxel_size):
        pcd_down = pcd.voxel_down_sample(voxel_size)
        pcd_down.estimate_normals(
            o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
        # Apply statistical outlier removal
        cl, ind = pcd_down.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
        pcd_down = pcd_down.select_by_index(ind)
        return pcd_down

    def pairwise_registration(source, target, use_colored_icp, voxel_size, max_correspondence_distance_coarse, max_correspondence_distance_fine):
        current_transformation = np.identity(4)  # Start with identity matrix

        if use_colored_icp:
            print("Apply colored point cloud registration")
            voxel_radius = [5*voxel_size, 3*voxel_size, voxel_size]
            max_iter = [60, 35, 20]
            
            for scale in range(3):
                iter = max_iter[scale]
                radius = voxel_radius[scale]

                source_down = source.voxel_down_sample(radius)
                target_down = target.voxel_down_sample(radius)

                source_down.estimate_normals(
                    o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))
                target_down.estimate_normals(
                    o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))

                try:
                    result_icp = o3d.pipelines.registration.registration_colored_icp(
                        source_down, target_down, radius, current_transformation,
                        o3d.pipelines.registration.TransformationEstimationForColoredICP(),
                        o3d.pipelines.registration.ICPConvergenceCriteria(relative_fitness=1e-6,
                                                                        relative_rmse=1e-6,
                                                                        max_iteration=iter))
                    current_transformation = result_icp.transformation
                except RuntimeError as e:
                    print(f"Colored ICP failed at scale {scale}: {str(e)}")
                    print("Keeping the previous transformation")
                    # We keep the previous transformation, no need to reassign

            transformation_icp = current_transformation
        else:
            print("Apply point-to-plane ICP")
            try:
                icp_coarse = o3d.pipelines.registration.registration_icp(
                    source, target, max_correspondence_distance_coarse, current_transformation,
                    o3d.pipelines.registration.TransformationEstimationPointToPlane())
                current_transformation = icp_coarse.transformation
                
                icp_fine = o3d.pipelines.registration.registration_icp(
                    source, target, max_correspondence_distance_fine,
                    current_transformation,
                    o3d.pipelines.registration.TransformationEstimationPointToPlane())
                transformation_icp = icp_fine.transformation
            except RuntimeError as e:
                print(f"Point-to-plane ICP failed: {str(e)}")
                print("Keeping the best available transformation")
                transformation_icp = current_transformation

        try:
            information_icp = o3d.pipelines.registration.get_information_matrix_from_point_clouds(
                source, target, max_correspondence_distance_fine,
                transformation_icp)
        except RuntimeError as e:
            print(f"Failed to compute information matrix: {str(e)}")
            print("Using identity information matrix")
            information_icp = np.identity(6)

        return transformation_icp, information_icp

    def full_registration(pcds_down):
        pose_graph = o3d.pipelines.registration.PoseGraph()
        odometry = np.identity(4)
        pose_graph.nodes.append(o3d.pipelines.registration.PoseGraphNode(odometry))
        n_pcds = len(pcds_down)
        
        pairs = []
        for i in range(n_pcds - 1):
            for j in range(i + 1, min(i + overlap + 1, n_pcds)):
                pairs.append((i, j))
                if quadratic_overlap:
                    q = 2**(j-i)
                    if q > overlap and i + q < n_pcds:
                        pairs.append((i, i + q))
        
        for source_id, target_id in pairs:
            transformation_icp, information_icp = pairwise_registration(
                pcds_down[source_id], pcds_down[target_id], use_colored_icp,
                voxel_size, max_correspondence_distance_coarse, max_correspondence_distance_fine)
            print(f"Build PoseGraph: {source_id} -> {target_id}")
            
            if target_id == source_id + 1:
                odometry = np.dot(transformation_icp, odometry)
                pose_graph.nodes.append(
                    o3d.pipelines.registration.PoseGraphNode(
                        np.linalg.inv(odometry)))
            
            pose_graph.edges.append(
                o3d.pipelines.registration.PoseGraphEdge(source_id,
                                                         target_id,
                                                         transformation_icp,
                                                         information_icp,
                                                         uncertain=False))
        return pose_graph

    # Preprocess point clouds
    print("Preprocessing point clouds...")
    pcds_down = [preprocess_point_cloud(pcd, voxel_size) for pcd in pcds]

    print("Full registration ...")
    pose_graph = full_registration(pcds_down)
    
    print("Optimizing PoseGraph ...")
    option = o3d.pipelines.registration.GlobalOptimizationOption(
        max_correspondence_distance=max_correspondence_distance_fine,
        edge_prune_threshold=0.25,
        reference_node=0)
    
    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
        o3d.pipelines.registration.global_optimization(
            pose_graph,
            o3d.pipelines.registration.GlobalOptimizationLevenbergMarquardt(),
            o3d.pipelines.registration.GlobalOptimizationConvergenceCriteria(),
            option)

    print("Transform points and combine")
    pcd_combined = o3d.geometry.PointCloud()
    for point_id in range(len(pcds)):
        print(pose_graph.nodes[point_id].pose)
        pcds[point_id].transform(pose_graph.nodes[point_id].pose)
        pcd_combined += pcds[point_id]
    
    return pcd_combined, pose_graph