import numpy as np
import open3d as o3d
import torch

class PtsUtil:

    @staticmethod
    def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005):
        o3d_pc = o3d.geometry.PointCloud()
        o3d_pc.points = o3d.utility.Vector3dVector(point_cloud)
        downsampled_pc = o3d_pc.voxel_down_sample(voxel_size)
        return np.asarray(downsampled_pc.points)
    
    @staticmethod
    def transform_point_cloud(points, pose_mat):
        points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)
        points_h = np.dot(pose_mat, points_h.T).T
        return points_h[:, :3]
    
    @staticmethod
    def random_downsample_point_cloud(point_cloud, num_points):
        idx = np.random.choice(len(point_cloud), num_points, replace=True)
        return point_cloud[idx]
    
    @staticmethod
    def random_downsample_point_cloud_tensor(point_cloud, num_points):
        idx = torch.randint(0, len(point_cloud), (num_points,))
        return point_cloud[idx]