import numpy as np
import open3d as o3d


def _draw_pointcloud(pcd, T):

    R = T[0:3, 0:3]
    point = T[0:3, 3:4].flatten()
    x, y, z = R[:, 0], R[:, 1], R[:, 2]

    points = [
        [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
        point, point + x, point + y, point + z

    ]  # 画点：原点、第一主成分、第二主成分
    lines = [
        [0, 1], [0, 2], [0, 3],
        [4, 5], [4, 6], [4, 7]
    ]  # 画出三点之间两两连线
    colors = [
        [1, 0, 0], [0, 1, 0], [0, 0, 1],
        [1, 0, 0], [0, 1, 0], [0, 0, 1]
    ]

    line_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(points), lines=o3d.utility.Vector2iVector(lines))
    line_set.colors = o3d.utility.Vector3dVector(colors)

    pcd.append(line_set)
    o3d.visualization.draw_geometries(pcd)


def _preprocess_point_cloud(pcd, voxel_size):

    pcd_down = pcd.voxel_down_sample(voxel_size)
    radius_normal = voxel_size * 2

    pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))

    radius_feature = voxel_size * 5

    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100)
    )

    return pcd_down, pcd_fpfh

def _prepare_dataset(source, target, voxel_size):

    trans_init = np.identity(4)
    source.transform(trans_init)

    source_down, source_fpfh = _preprocess_point_cloud(source, voxel_size)
    target_down, target_fpfh = _preprocess_point_cloud(target, voxel_size)

    return 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 * 1.5
    result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        source_down,
        target_down,
        source_fpfh,
        target_fpfh,
        True,
        distance_threshold,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
        3,
        [
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold)
        ],
        o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999))

    return result.transformation


def object_icp(
        target: o3d.geometry.PointCloud,
        source: o3d.geometry.PointCloud | str,
        ransac_voxel_size: float = 0.005,
        icp_voxel_radius: list[float] | tuple[float] | None = (0.004, 0.002, 0.001),
        icp_max_iter: list[int] | tuple[int] | None = (50, 30, 14),
):
    if isinstance(source, str):
        source = o3d.io.read_point_cloud(source)
    elif isinstance(source, o3d.geometry.PointCloud):
        pass
    else:
        raise TypeError(f"Unsupported Type {type(source)}")

    voxel_size = 0.005  # means 5mm for this dataset
    source_down, target_down, source_fpfh, target_fpfh = _prepare_dataset(source, target, voxel_size)

    T = _execute_global_registration(source_down, target_down, source_fpfh, target_fpfh, ransac_voxel_size)

    for scale in range(3):
        iter = icp_max_iter[scale]
        radius = icp_voxel_radius[scale]
        # print([iter, 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))

        result_icp = o3d.pipelines.registration.registration_icp(
            source_down,
            target_down,
            radius * 5,
            T,
            o3d.pipelines.registration.TransformationEstimationPointToPoint(),
            o3d.pipelines.registration.ICPConvergenceCriteria(
                relative_fitness=1e-6,
                relative_rmse=1e-6,
                max_iteration=iter
            )
        )

        T = result_icp.transformation

    _draw_pointcloud([source.transform(T), target], T)
    
    return T


if __name__ == "__main__":
    target = o3d.io.read_point_cloud("pointcloud/pointcloud_0.pcd")
    source = o3d.io.read_point_cloud("pointcloud/bottle_model.pcd")
    Rmat = object_icp(target, source)


