point_cloud_np = np.loadtxt("C:\\Users\\jtz08\\PycharmProjects\\untitled\\airplane_0544.txt", delimiter=',')
point_cloud_np = point_cloud_np[:, 0:3]
point_cloud_o3d = o3d.geometry.PointCloud()
point_cloud_o3d.points = o3d.utility.Vector3dVector(point_cloud_np)

o3d.visualization.draw_geometries([point_cloud_o3d])

def PCA(points, correlation=False, sort=True):
    points_mean = np.mean(points, axis=0)
    points_normalized = points - points_mean
    func = np.cov if not correlation else np.corrcoef
    H = func(points_normalized, rowvar=False, bias=True)
    vectors, values, vectors_t = np.linalg.svd(H, full_matrices=True)
    if sort:
        sort = values.argsort()[::-1]
        values = values[sort]
        vectors = vectors[:, sort]
    return values, vectors

points = np.asarray(point_cloud_o3d.points)
values, vectors = PCA(points)
point_cloud_vector = vectors[:, 0]

def PCA_3d(points, vectors):
    points_mean = np.mean(points, axis=0)
    pointset = [points_mean,
             points_mean + vectors[:, 0], points_mean - vectors[:, 0],
             points_mean + vectors[:, 1], points_mean - vectors[:, 1],
             points_mean + vectors[:, 2], points_mean - vectors[:, 2]]
    lines = [[0, 1], [0, 2], [0, 3], [0, 4], [0, 5], [0, 6]]
    colors = [[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1]]
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(pointset)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector(colors)
    return line_set

line_set = PCA_3d(points, vectors)
o3d.visualization.draw_geometries([point_cloud_o3d, line_set])

def PCA_2d(points, vectors):
    projected_points = np.dot(points, vectors[:, 0:2])
    zeros = np.zeros((projected_points.shape[0], 1))
    projected_points = np.hstack((projected_points, zeros))
    point_cloud_projected = o3d.geometry.PointCloud()
    point_cloud_projected.points = o3d.utility.Vector3dVector(projected_points)
    return point_cloud_projected

def PCA_2d_23(point_cloud_o3d, vectors):
    X = np.asarray(point_cloud_o3d.points)
    projected_points = np.dot(X, vectors[:, 1:])
    zeros = np.zeros((projected_points.shape[0], 1))
    projected_points = np.hstack((projected_points, zeros))
    point_cloud_projected = o3d.geometry.PointCloud()
    point_cloud_projected.points = o3d.utility.Vector3dVector(projected_points)
    return point_cloud_projected

def PCA_2d_13(point_cloud_o3d, vectors):
    X = np.asarray(point_cloud_o3d.points)
    projected_points = np.dot(X, vectors.take([0,2], 1))
    zeros = np.zeros((projected_points.shape[0], 1))
    projected_points = np.hstack((projected_points, zeros))
    point_cloud_projected = o3d.geometry.PointCloud()
    point_cloud_projected.points = o3d.utility.Vector3dVector(projected_points)
    return point_cloud_projected

point_cloud_projected = PCA_2d(points, vectors)
point_cloud_projected_23 = PCA_2d_23(point_cloud_o3d, vectors)
point_cloud_projected_13 = PCA_2d_13(point_cloud_o3d, vectors)

o3d.visualization.draw_geometries([point_cloud_projected])
o3d.visualization.draw_geometries([point_cloud_projected_23])
o3d.visualization.draw_geometries([point_cloud_projected_13])

def normals(point_cloud_o3d):
    pcd_tree = o3d.geometry.KDTreeFlann(point_cloud_o3d)
    points = np.asarray(point_cloud_o3d.points)
    normals = []
    for point in points:
        k, idx, _ = pcd_tree.search_knn_vector_3d(point, knn=50)
        value, vector = PCA(points[idx, :])
        normals.append(vector[:, -1])
    return np.array(normals, dtype=np.float64)

normals = normals(point_cloud_o3d)
point_cloud_o3d.normals = o3d.utility.Vector3dVector(normals)
o3d.visualization.draw_geometries([point_cloud_o3d])

def normals_o3d(points, normals, scale=0.08):
    N = points.shape[0]
    pointset = np.vstack((points, points + normals*scale))
    lines = [[i, i+N] for i in range(N)]
    colors = np.array([[0.58, 0, 0.83]*N]).reshape(N, 3)
    normals_o3d = o3d.geometry.LineSet()
    normals_o3d.points = o3d.utility.Vector3dVector(pointset)
    normals_o3d.lines = o3d.utility.Vector2iVector(lines)
    normals_o3d.colors = o3d.utility.Vector3dVector(colors)
    return normals_o3d

normals_o3d = normals_o3d(points, normals)
o3d.visualization.draw_geometries([point_cloud_o3d, normals_o3d])

def voxel_fileter(points, leaf_size, method='random'):
    xyz_max = np.max(points, axis=0)
    xyz_min = np.min(points, axis=0)
    D_xyz = np.ceil((xyz_max - xyz_min) / leaf_size)
    idx = np.floor((points - xyz_min) / leaf_size)
    h = idx[:, 0] + idx[:, 1] * D_xyz[0] + idx[:, 2] * D_xyz[0] * D_xyz[1]
    data = np.c_[h, points]
    data = data[data[:, 0].argsort()]
    filtered_points = []
    if method == 'random':
        for i in range(data.shape[0] - 1):
            if data[i][0] != data[i+1][0]:
                filtered_points.append(data[i][1:])
        filtered_points.append(data[data.shape[0] - 1][1:])
        return np.array(filtered_points, dtype=np.float64)
    elif method == 'centroid':
        data_points = []
        for i in range(data.shape[0] - 1):
            if data[i][0] == data[i+1][0]:
                data_points.append(data[i][1:])
                continue
            elif data_points == []:
                continue
                # filtered_points.append(data[i][1:])
            else:
                data_points.append(data[i][1:])
                filtered_points.append(np.mean(data_points, axis=0))
                data_points = []
        return np.array(filtered_points, dtype=np.float64)

filtered_points_random = voxel_fileter(points, 0.06, method='random')
filtered_cloud_o3d_random = o3d.geometry.PointCloud()
filtered_cloud_o3d_random.points = o3d.utility.Vector3dVector(filtered_points_random)
o3d.visualization.draw_geometries([filtered_cloud_o3d_random])

filtered_points_centroid = voxel_fileter(points, 0.06, method='centroid')
filtered_cloud_o3d_centroid = o3d.geometry.PointCloud()
filtered_cloud_o3d_centroid.points = o3d.utility.Vector3dVector(filtered_points_centroid)
o3d.visualization.draw_geometries([filtered_cloud_o3d_centroid])