import copy

import numpy as np
import open3d as o3d
from tqdm import tqdm
from scipy import stats

import utils_o3d as utils


def remove_ground_plane(pcd, z_thresh=-2.7):
    cropped = copy.deepcopy(pcd)
    cropped_points = np.array(cropped.points)
    cropped_points = cropped_points[cropped_points[:, -1] > z_thresh]
    pcd_final = o3d.geometry.PointCloud()
    pcd_final.points = o3d.utility.Vector3dVector(cropped_points)
    return pcd_final


def remove_y_plane(pcd, y_thresh=5):
    cropped = copy.deepcopy(pcd)
    cropped_points = np.array(cropped.points)
    cropped_points = cropped_points[cropped_points[:, 0] < y_thresh]
    cropped_points[:, -1] = -cropped_points[:, -1]
    pcd_final = o3d.geometry.PointCloud()
    pcd_final.points = o3d.utility.Vector3dVector(cropped_points)
    return pcd_final


def compute_features(pcd, voxel_size, normals_nn=100, features_nn=120, downsample=True):
    normals_radius  = voxel_size * 2
    features_radius = voxel_size * 4

    # Downsample the point cloud using Voxel grids 
    if downsample:
        print(':: Input size:', np.array(pcd.points).shape)
        pcd_down = utils.downsample_point_cloud(pcd, voxel_size)
        print(':: Downsample with a voxel size %.3f' % voxel_size)
        print(':: Downsample size', np.array(pcd_down.points).shape)
    else: pcd_down = copy.deepcopy(pcd)

    # Estimate normals  
    print(':: Estimate normal with search radius %.3f' % normals_radius)
    pcd_down.estimate_normals(
        o3d.geometry.KDTreeSearchParamHybrid(radius=normals_radius, max_nn=normals_nn))

    # Compute FPFH features 
    print(':: Compute FPFH feature with search radius %.3f' % features_radius)
    features = o3d.registration.compute_fpfh_feature(pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=features_radius, max_nn=features_nn))

    return pcd_down, features


def match_features(pcd0, pcd1, feature0, feature1, thresh=None, display=False):
    pcd0, pcd1 = copy.deepcopy(pcd0), copy.deepcopy(pcd1)
    print(':: Input size 0:', np.array(pcd0.points).shape)
    print(':: Input size 1:', np.array(pcd1.points).shape)
    print(':: Features size 0:', np.array(feature0.data).shape)
    print(':: Features size 1:', np.array(feature1.data).shape)

    utils.paint_uniform_color(pcd0, color=[1, 0.706, 0])
    utils.paint_uniform_color(pcd1, color=[0, 0.651, 0.929])

    scores, indices = [], []
    fpfh_tree = o3d.geometry.KDTreeFlann(feature1)
    for i in tqdm(range(len(pcd0.points)), desc=':: Feature Matching'):
        [_, idx, _] = fpfh_tree.search_knn_vector_xd(feature0.data[:, i], 1)
        scores.append(np.linalg.norm(pcd0.points[i] - pcd1.points[idx[0]]))
        indices.append([i, idx[0]])

    scores, indices = np.array(scores), np.array(indices)
    median = np.median(scores)
    if thresh is None: thresh = median
    inliers_idx = np.where(scores <= thresh)[0]
    pcd0_idx = indices[inliers_idx, 0]
    pcd1_idx = indices[inliers_idx, 1]

    print(':: Score stats: Min=%0.3f, Max=%0.3f, Median=%0.3f, N<Thresh=%d' % (
        np.min(scores), np.max(scores), median, len(inliers_idx)))

    if display:
        for i, j in zip(pcd0_idx, pcd1_idx):
            pcd0.colors[i] = [1, 0, 0]
            pcd1.colors[j] = [1, 0, 0]
        utils.display([pcd0, pcd1])

    return pcd0_idx, pcd1_idx


def estimate_scale(pcd0, pcd1, pcd0_idx, pcd1_idx, top_percent=1.0,
    ransac_iters=5000, sample_size=50):
    points0 = np.asarray(pcd0.points)[pcd0_idx]
    points1 = np.asarray(pcd1.points)[pcd1_idx]
    mean0 = np.mean(points0, axis=0)
    mean1 = np.mean(points1, axis=0)
    top_count = int(top_percent * len(pcd0_idx))
    assert top_count > sample_size, 'top_count <= sample_size'

    scales = []
    for i in tqdm(range(ransac_iters), desc=':: Scale Estimation RANSAC'):
        args = np.random.choice(top_count, sample_size, replace=False)
        points0_r = points0[args]
        points1_r = points1[args]

        score0 = np.sum((points0_r - mean0) ** 2, axis=1)
        score1 = np.sum((points1_r - mean1) ** 2, axis=1)
        scale = np.sqrt(np.mean(score1) / np.mean(score0))
        scales.append(scale)

    best_scale = stats.mode(scales)[0][0]
    print(':: Estimated scale:', best_scale)

    return best_scale


def global_registration(source_down, target_down, source_fpfh, target_fpfh, voxel_size,
    distance_threshold=1.0, num_iters=4000000, num_val_iters=500):
    print(':: 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(num_iters, num_val_iters))

    return result


def fast_global_registration(source_down, target_down,
    source_fpfh, target_fpfh, voxel_size):
    distance_threshold = 1.0
    result = o3d.registration.registration_fast_based_on_feature_matching(
        source_down, target_down, source_fpfh, target_fpfh,
        o3d.registration.FastGlobalRegistrationOption(
            maximum_correspondence_distance=distance_threshold))

    return result


def refine_registration(source, target, source_fpfh, target_fpfh, initial_result, voxel_size):
    distance_threshold = 0.1
    print(':: Distance threshold %.3f' % distance_threshold)

    result = o3d.registration.registration_icp(
        source, target, distance_threshold, initial_result.transformation,
        o3d.registration.TransformationEstimationPointToPlane())

    return result


def registration(pcd0, pcd1, feature1, feature2, voxel_size, method='global'):
    if method == 'global':
        print('\nRANSAC global registration on scaled point clouds...')
        initial_result = global_registration(pcd0, pcd1, feature1, feature2, voxel_size)

    elif method == 'fast_global':
        print('\nFast global registration on scaled point clouds...')
        initial_result = fast_global_registration(pcd0, pcd1, feature1, feature2, voxel_size)

    else:
        print(':: Registration method not supported')
        return

    print(':: Initial registration results:')
    print(initial_result)

    print('\nDisplaying initial result...')
    draw_registration_result(pcd0, pcd1, initial_result.transformation)

    print('\nRefine registration...')
    result = refine_registration(pcd0, pcd1, feature1, feature2, initial_result, voxel_size)

    print(':: Final registration results:')
    print(result)

    return result


def draw_registration_result(source, target, transformation):
    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)
    o3d.visualization.draw_geometries([source_temp, target_temp])


def run():
    voxel_size = 0.2
    dso_scale = 0.03

    pcd_lidar = o3d.io.read_point_cloud('../maps/scans/scan_050.pcd')
    pcd_lidar = remove_ground_plane(pcd_lidar)

    pcd_dso = o3d.io.read_point_cloud('../maps/dso_map_cleaned.pcd')
    pcd_dso = remove_ground_plane(pcd_dso, z_thresh=4.5)
    pcd_dso = remove_y_plane(pcd_dso, y_thresh=0.2)
    # pcd_dso = utils.scale_point_cloud(pcd_dso, dso_scale).rotate([0.5, 0.5, 0.5]).translate([10, 20, 30])

    # Ground plane removal results
    # utils.display(pcds=[pcd_lidar, pcd_dso], colors=[[1, 0.706, 0], [0, 0.651, 0.929]])
    # utils.display(pcds=[pcd_dso], colors=[[0, 0.651, 0.929]])
    # return

    print('\nComputing FPFH features for lidar point cloud...')
    pcd_lidar_down, features_lidar = compute_features(pcd_lidar, voxel_size=voxel_size)

    print('\nComputing FPFH features for DSO point cloud...')
    pcd_dso_down, features_dso = compute_features(pcd_dso, voxel_size=voxel_size * (dso_scale if dso_scale < 1 else 1))

    print('\nMatching FPFH features...')
    pcd_lidar_idx, pcd_dso_idx = match_features(pcd_lidar_down, pcd_dso_down,
        features_lidar, features_dso, thresh=None)

    print('\nEstimating scale using matches...')
    scale = estimate_scale(pcd_lidar_down, pcd_dso_down, pcd_lidar_idx, pcd_dso_idx)
    scale = 0.06

    print('\nCorrecting scale...')
    pcd_dso_scaled = utils.scale_point_cloud(pcd_dso, 1.0 / scale)
    utils.display(pcds=[pcd_lidar, pcd_dso_scaled], colors=[[1, 0.706, 0], [0, 0.651, 0.929]])
    # return

    # Registration
    pcd_dso_scaled_down, features_dso_scaled = compute_features(
        pcd_dso_scaled, voxel_size=voxel_size)
    result = registration(pcd_lidar_down, pcd_dso_scaled_down, features_lidar,
        features_dso_scaled, voxel_size, method='global')

    print('\nDisplaying result...')
    draw_registration_result(pcd_lidar, pcd_dso_scaled, result.transformation)


if __name__ == '__main__':
    run()    

[((12, 14, 12, 32), 'copy.deepcopy', 'copy.deepcopy', ({(12, 28, 12, 31): 'pcd'}, {}), '(pcd)', False, 'import copy\n'), ((13, 21, 13, 45), 'numpy.array', 'np.array', ({(13, 30, 13, 44): 'cropped.points'}, {}), '(cropped.points)', True, 'import numpy as np\n'), ((15, 16, 15, 41), 'open3d.geometry.PointCloud', 'o3d.geometry.PointCloud', ({}, {}), '()', True, 'import open3d as o3d\n'), ((16, 23, 16, 65), 'open3d.utility.Vector3dVector', 'o3d.utility.Vector3dVector', ({(16, 50, 16, 64): 'cropped_points'}, {}), '(cropped_points)', True, 'import open3d as o3d\n'), ((21, 14, 21, 32), 'copy.deepcopy', 'copy.deepcopy', ({(21, 28, 21, 31): 'pcd'}, {}), '(pcd)', False, 'import copy\n'), ((22, 21, 22, 45), 'numpy.array', 'np.array', ({(22, 30, 22, 44): 'cropped.points'}, {}), '(cropped.points)', True, 'import numpy as np\n'), ((25, 16, 25, 41), 'open3d.geometry.PointCloud', 'o3d.geometry.PointCloud', ({}, {}), '()', True, 'import open3d as o3d\n'), ((26, 23, 26, 65), 'open3d.utility.Vector3dVector', 'o3d.utility.Vector3dVector', ({(26, 50, 26, 64): 'cropped_points'}, {}), '(cropped_points)', True, 'import open3d as o3d\n'), ((62, 4, 62, 56), 'utils_o3d.paint_uniform_color', 'utils.paint_uniform_color', (), '', True, 'import utils_o3d as utils\n'), ((63, 4, 63, 60), 'utils_o3d.paint_uniform_color', 'utils.paint_uniform_color', (), '', True, 'import utils_o3d as utils\n'), ((66, 16, 66, 50), 'open3d.geometry.KDTreeFlann', 'o3d.geometry.KDTreeFlann', ({(66, 41, 66, 49): 'feature1'}, {}), '(feature1)', True, 'import open3d as o3d\n'), ((73, 13, 73, 30), 'numpy.median', 'np.median', ({(73, 23, 73, 29): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((95, 12, 95, 36), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((96, 12, 96, 36), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((183, 18, 183, 39), 'copy.deepcopy', 'copy.deepcopy', ({(183, 32, 183, 38): 'source'}, {}), '(source)', False, 'import copy\n'), ((184, 18, 184, 39), 'copy.deepcopy', 'copy.deepcopy', ({(184, 32, 184, 38): 'target'}, {}), '(target)', False, 'import copy\n'), ((188, 4, 188, 65), 'open3d.visualization.draw_geometries', 'o3d.visualization.draw_geometries', ({(188, 38, 188, 64): '[source_temp, target_temp]'}, {}), '([source_temp, target_temp])', True, 'import open3d as o3d\n'), ((195, 16, 195, 69), 'open3d.io.read_point_cloud', 'o3d.io.read_point_cloud', ({(195, 40, 195, 68): '"""../maps/scans/scan_050.pcd"""'}, {}), "('../maps/scans/scan_050.pcd')", True, 'import open3d as o3d\n'), ((198, 14, 198, 68), 'open3d.io.read_point_cloud', 'o3d.io.read_point_cloud', ({(198, 38, 198, 67): '"""../maps/dso_map_cleaned.pcd"""'}, {}), "('../maps/dso_map_cleaned.pcd')", True, 'import open3d as o3d\n'), ((223, 21, 223, 66), 'utils_o3d.scale_point_cloud', 'utils.scale_point_cloud', ({(223, 45, 223, 52): 'pcd_dso', (223, 54, 223, 65): '1.0 / scale'}, {}), '(pcd_dso, 1.0 / scale)', True, 'import utils_o3d as utils\n'), ((224, 4, 224, 94), 'utils_o3d.display', 'utils.display', (), '', True, 'import utils_o3d as utils\n'), ((37, 19, 37, 64), 'utils_o3d.downsample_point_cloud', 'utils.downsample_point_cloud', ({(37, 48, 37, 51): 'pcd', (37, 53, 37, 63): 'voxel_size'}, {}), '(pcd, voxel_size)', True, 'import utils_o3d as utils\n'), ((40, 21, 40, 39), 'copy.deepcopy', 'copy.deepcopy', ({(40, 35, 40, 38): 'pcd'}, {}), '(pcd)', False, 'import copy\n'), ((45, 8, 45, 86), 'open3d.geometry.KDTreeSearchParamHybrid', 'o3d.geometry.KDTreeSearchParamHybrid', (), '', True, 'import open3d as o3d\n'), ((50, 8, 50, 88), 'open3d.geometry.KDTreeSearchParamHybrid', 'o3d.geometry.KDTreeSearchParamHybrid', (), '', True, 'import open3d as o3d\n'), ((56, 17, 56, 36), 'copy.deepcopy', 'copy.deepcopy', ({(56, 31, 56, 35): 'pcd0'}, {}), '(pcd0)', False, 'import copy\n'), ((56, 38, 56, 57), 'copy.deepcopy', 'copy.deepcopy', ({(56, 52, 56, 56): 'pcd1'}, {}), '(pcd1)', False, 'import copy\n'), ((72, 22, 72, 38), 'numpy.array', 'np.array', ({(72, 31, 72, 37): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((72, 40, 72, 57), 'numpy.array', 'np.array', ({(72, 49, 72, 56): 'indices'}, {}), '(indices)', True, 'import numpy as np\n'), ((75, 18, 75, 44), 'numpy.where', 'np.where', ({(75, 27, 75, 43): '(scores <= thresh)'}, {}), '(scores <= thresh)', True, 'import numpy as np\n'), ((86, 8, 86, 35), 'utils_o3d.display', 'utils.display', ({(86, 22, 86, 34): '[pcd0, pcd1]'}, {}), '([pcd0, pcd1])', True, 'import utils_o3d as utils\n'), ((93, 14, 93, 37), 'numpy.asarray', 'np.asarray', ({(93, 25, 93, 36): 'pcd0.points'}, {}), '(pcd0.points)', True, 'import numpy as np\n'), ((94, 14, 94, 37), 'numpy.asarray', 'np.asarray', ({(94, 25, 94, 36): 'pcd1.points'}, {}), '(pcd1.points)', True, 'import numpy as np\n'), ((102, 15, 102, 70), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((106, 17, 106, 57), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((107, 17, 107, 57), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((123, 8, 123, 68), 'open3d.registration.TransformationEstimationPointToPoint', 'o3d.registration.TransformationEstimationPointToPoint', ({(123, 62, 123, 67): 'False'}, {}), '(False)', True, 'import open3d as o3d\n'), ((127, 11, 127, 79), 'open3d.registration.RANSACConvergenceCriteria', 'o3d.registration.RANSACConvergenceCriteria', ({(127, 54, 127, 63): 'num_iters', (127, 65, 127, 78): 'num_val_iters'}, {}), '(num_iters, num_val_iters)', True, 'import open3d as o3d\n'), ((137, 8, 138, 63), 'open3d.registration.FastGlobalRegistrationOption', 'o3d.registration.FastGlobalRegistrationOption', (), '', True, 'import open3d as o3d\n'), ((149, 8, 149, 63), 'open3d.registration.TransformationEstimationPointToPlane', 'o3d.registration.TransformationEstimationPointToPlane', ({}, {}), '()', True, 'import open3d as o3d\n'), ((57, 30, 57, 51), 'numpy.array', 'np.array', ({(57, 39, 57, 50): 'pcd0.points'}, {}), '(pcd0.points)', True, 'import numpy as np\n'), ((58, 30, 58, 51), 'numpy.array', 'np.array', ({(58, 39, 58, 50): 'pcd1.points'}, {}), '(pcd1.points)', True, 'import numpy as np\n'), ((59, 33, 59, 56), 'numpy.array', 'np.array', ({(59, 42, 59, 55): 'feature0.data'}, {}), '(feature0.data)', True, 'import numpy as np\n'), ((60, 33, 60, 56), 'numpy.array', 'np.array', ({(60, 42, 60, 55): 'feature1.data'}, {}), '(feature1.data)', True, 'import numpy as np\n'), ((69, 22, 69, 74), 'numpy.linalg.norm', 'np.linalg.norm', ({(69, 37, 69, 73): '(pcd0.points[i] - pcd1.points[idx[0]])'}, {}), '(pcd0.points[i] - pcd1.points[idx[0]])', True, 'import numpy as np\n'), ((111, 17, 111, 35), 'scipy.stats.mode', 'stats.mode', ({(111, 28, 111, 34): 'scales'}, {}), '(scales)', False, 'from scipy import stats\n'), ((124, 12, 124, 72), 'open3d.registration.CorrespondenceCheckerBasedOnEdgeLength', 'o3d.registration.CorrespondenceCheckerBasedOnEdgeLength', ({(124, 68, 124, 71): '0.9'}, {}), '(0.9)', True, 'import open3d as o3d\n'), ((125, 12, 126, 35), 'open3d.registration.CorrespondenceCheckerBasedOnDistance', 'o3d.registration.CorrespondenceCheckerBasedOnDistance', ({(126, 16, 126, 34): 'distance_threshold'}, {}), '(distance_threshold)', True, 'import open3d as o3d\n'), ((36, 32, 36, 52), 'numpy.array', 'np.array', ({(36, 41, 36, 51): 'pcd.points'}, {}), '(pcd.points)', True, 'import numpy as np\n'), ((39, 36, 39, 61), 'numpy.array', 'np.array', ({(39, 45, 39, 60): 'pcd_down.points'}, {}), '(pcd_down.points)', True, 'import numpy as np\n'), ((80, 8, 80, 22), 'numpy.min', 'np.min', ({(80, 15, 80, 21): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((80, 24, 80, 38), 'numpy.max', 'np.max', ({(80, 31, 80, 37): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((108, 24, 108, 39), 'numpy.mean', 'np.mean', ({(108, 32, 108, 38): 'score1'}, {}), '(score1)', True, 'import numpy as np\n'), ((108, 42, 108, 57), 'numpy.mean', 'np.mean', ({(108, 50, 108, 56): 'score0'}, {}), '(score0)', True, 'import numpy as np\n')]