import open3d as o3d
import numpy as np
import copy
from datetime import datetime

count = 0
corr_treshold = 1000

def draw_registration_result(source, target, transformation, info):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    if info == 1:
        o3d.visualization.draw_geometries([source_temp, target_temp], window_name="RANSAC")

    if info == 2:
        o3d.visualization.draw_geometries([source_temp, target_temp], window_name="ICP")

    if info == 0:
        o3d.visualization.draw_geometries([source_temp, target_temp])




def preprocess_point_cloud(pcd, voxel_size):

    #print(":: Downsample with a voxel size %.3f." % voxel_size)
    pcldownsampled = pcd.voxel_down_sample(voxel_size)
    #print("Radius oulier removal")
    cl, ind = pcldownsampled.remove_radius_outlier(nb_points=16, radius=5)
    pclSampleFirst = pcldownsampled.select_down_sample(ind)
    #print(pclSampleFirst)
    #print("Statistic oulier removal")
    cl2, ind = pclSampleFirst.remove_statistical_outlier(nb_neighbors=20,
                                                         std_ratio=2.0)
    pcd_down = pclSampleFirst.select_down_sample(ind)
    #print(pcd_down)

    radius_normal = voxel_size * 4
    # print(":: Estimate normal with search radius %.3f." %radius_normal)
    pcd_down.estimate_normals(
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=50))

    radius_feature = voxel_size * 50
    # print(":: Compute FPFH feature with search radius %.3f." % radius_feature)
    pcd_fpfh = o3d.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=150))
    return pcd_down, pcd_fpfh


def prepare_dataset(voxel_size):
    # print(":: Load two point clouds and disturb initial pose.")
    source = o3d.io.read_point_cloud("test_sphere - Cloud.pcd")
    print(source)
    target = o3d.io.read_point_cloud("Sphere_r=60.ply")
    draw_registration_result(source, target, np.identity(4), 0)
    source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
    # print(source_down)
    target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)
    return source, target, source_down, target_down, source_fpfh, target_fpfh


def execute_global_registration(source_down, target_down, source_fpfh,
                                target_fpfh, voxel_size):
    distance_threshold = voxel_size * 150
    # print(":: RANSAC registration on downsampled point clouds.")
    # print("   Since the downsampling voxel size is %.3f," % voxel_size)
    # print("   we use a liberal distance threshold %.3f." % distance_threshold)
    result = o3d.registration.registration_ransac_based_on_feature_matching(
        source_down, target_down, source_fpfh, target_fpfh, distance_threshold,
        o3d.registration.TransformationEstimationPointToPoint(False), 4, [
            o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.registration.CorrespondenceCheckerBasedOnDistance(
                distance_threshold)
        ], o3d.registration.RANSACConvergenceCriteria(900000000, 100000))
    if len(np.asarray(result.correspondence_set)) < corr_treshold:
        print(len(np.asarray(result.correspondence_set)))
        print("TEMPLATE NOT FOUND IN SOURCE")
        result = 0
        print(result)
    else:
        print(len(np.asarray(result.correspondence_set)))
        print("Transformation is:")
        print(result.transformation)
        print("")
    return result


if __name__ == "__main__":
    start_time = datetime.now()
    voxel_size = 0.05   # means 5cm for the dataset
    source, target, source_down, target_down, source_fpfh, target_fpfh = \
         prepare_dataset(voxel_size)

    result_ransac = execute_global_registration(source_down, target_down,
                                                    source_fpfh, target_fpfh,
                                                    voxel_size)
    #print(result_ransac)
    if result_ransac != 0:
        time_elapsed = datetime.now() - start_time
        draw_registration_result(source_down, target_down,
                                result_ransac.transformation, 1)
        print("------INFO-------- ")
        print(len(np.asarray(result_ransac.correspondence_set)))

    else:
        print("Could not find properly")
        time_elapsed = datetime.now() - start_time

    print('Time elapsed {}'.format(time_elapsed))

