import os
import numpy as np
import trimesh
from utils.pts_util import PtsUtil
import numpy as np
import open3d as o3d
import torch
import trimesh
from scipy.spatial import cKDTree

def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.002):
    radius_normal = voxel_size * 2
    pipreg = o3d.pipelines.registration
    model_pcd = o3d.geometry.PointCloud()
    model_pcd.points = o3d.utility.Vector3dVector(model.vertices)
    model_downsampled = model_pcd.voxel_down_sample(voxel_size)
    model_downsampled.estimate_normals(
        search_param=o3d.geometry.KDTreeSearchParamHybrid(
            radius=radius_normal, max_nn=30
        )
    )
    # model_fpfh = pipreg.compute_fpfh_feature(
    #     model_downsampled,
    #     o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
    # )

    source_pcd = o3d.geometry.PointCloud()
    source_pcd.points = o3d.utility.Vector3dVector(pcl)
    source_downsampled = source_pcd.voxel_down_sample(voxel_size)
    source_downsampled.estimate_normals(
        search_param=o3d.geometry.KDTreeSearchParamHybrid(
            radius=radius_normal, max_nn=30
        )
    )
    # source_fpfh = pipreg.compute_fpfh_feature(
    #     source_downsampled,
    #     o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
    # )

    # reg_ransac = pipreg.registration_ransac_based_on_feature_matching(
    #     source_downsampled,
    #     model_downsampled,
    #     source_fpfh,
    #     model_fpfh,
    #     mutual_filter=True,
    #     max_correspondence_distance=voxel_size * 2,
    #     estimation_method=pipreg.TransformationEstimationPointToPoint(False),
    #     ransac_n=4,
    #     checkers=[pipreg.CorrespondenceCheckerBasedOnEdgeLength(0.9)],
    #     criteria=pipreg.RANSACConvergenceCriteria(4000000, 500),
    # )
    
    reg_icp2 = pipreg.registration_icp(
        source_downsampled,
        model_downsampled,
        voxel_size*10,
        np.eye(4),
        pipreg.TransformationEstimationPointToPlane(),
        pipreg.ICPConvergenceCriteria(max_iteration=2000),
    )
    
    reg_icp = pipreg.registration_icp(
        source_downsampled,
        model_downsampled,
        voxel_size,
        reg_icp2.transformation,
        pipreg.TransformationEstimationPointToPlane(),
        pipreg.ICPConvergenceCriteria(max_iteration=2000),
    )
    
    return reg_icp2.transformation, reg_icp.transformation

if __name__ == "__main__":

    model_dir = "/home/yan20/Desktop/nbv_rec/data/models/bear"
    model_path = os.path.join(model_dir,"mesh.ply")
    temp_name = "cad_model_world"
    cad_model = trimesh.load(model_path)
    
    pts_path = "/home/yan20/nbv_rec/project/franka_control/first_real_pts_bear.txt"
    pts = np.loadtxt(pts_path)
    very_coarse_real_to_cad = np.eye(4)
    cad_model_pts = cad_model.vertices
    very_coarse_real_to_cad[:3,3] = np.mean(cad_model_pts, axis=0) - np.mean(pts, axis=0)
    very_coarse_cad_pts = PtsUtil.transform_point_cloud(pts, very_coarse_real_to_cad)
    
    target_point = np.array([-3.422540776542676300e-02, -2.412379452948226755e-02, 1.123609126159126337e-01])

    # 设置一个容忍度
    tolerance = 1e-5

    # 计算每个点与目标点之间的距离
    distances = np.linalg.norm(very_coarse_cad_pts - target_point, axis=1)

    # 统计距离小于容忍度的点的数量
    count = np.sum(distances < tolerance)
    print(count)
    print(very_coarse_cad_pts.shape)
    print(np.mean(pts, axis=0), np.mean(cad_model_pts, axis=0), np.mean(very_coarse_cad_pts, axis=0))
    pts = pts[distances > tolerance]
    np.savetxt(os.path.join(temp_name + "_filtered.txt"), pts)
    # very_coarse_real_to_cad[:3,3] = np.mean(cad_model_pts, axis=0) - np.mean(pts, axis=0)
    # very_coarse_cad_pts = PtsUtil.transform_point_cloud(pts, very_coarse_real_to_cad)
    # np.savetxt(os.path.join(temp_name + "_very_coarse_reg.txt"), very_coarse_cad_pts)
    
    
    # real_to_cad = PtsUtil.register(very_coarse_cad_pts, cad_model)
    # cad_pts = PtsUtil.transform_point_cloud(very_coarse_cad_pts, real_to_cad)
    # np.savetxt(os.path.join(temp_name + "_reg.txt"), cad_pts)