import numpy as np
import open3d as o3d


# 读取文件
def open_file(file_path):
    """
    读取文件
    :param file_path: 文件路径
    :return: 读取的内容
    """

    if '.txt' in file_path:
        print('txt文档')
        pcd = o3d.io.read_point_cloud(file_path, format='xyz')
    else:
        pcd = o3d.io.read_point_cloud(file_path)
    return pcd


def save_file(content, file_path, write_ascii=True):
    """
    保存文件
    :param write_ascii: 是否使用ASCII编码
    :param content: 文件内容
    :param file_path: 保存路径
    :return: None
    """
    if '.txt' in file_path:
        np.savetxt(file_path, np.asarray(content.points), fmt='%f %f %f')
    else:
        o3d.io.write_point_cloud(file_path, content, write_ascii=write_ascii)


def voxel_down_sample(pcd, voxel_size):
    """
    体素方法点云降采样
    :param pcd: 点云数据
    :param voxel_size: 体素降采样的网格长度
    :return: 降采样点云数据
    """
    return pcd.voxel_down_sample(voxel_size=voxel_size)


def add_rgb_to_pcd(pcd):
    """
    给点云天啊及颜色信息
    :param pcd: 点云数据
    :return:
    """
    # 为点云中的每个点指定RGB颜色，这里我们使用红色作为示例
    # RGB值范围为[0, 1]，红色可以表示为(1, 0, 0)
    pcd = np.asarray(pcd.points).shape[0]
    colors = np.repeat([[1, 0, 0]], pcd, axis=0)  # 创建一个全是红色的颜色数组
    # 将颜色信息添加到点云对象中
    pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd


def pcd_visualization(pcd):
    """
    可视化点云数据
    :param pcd: 点云数据
    :return: None
    """
    o3d.visualization.draw_geometries([pcd],
                                      zoom=0.3412,
                                      front=[0.4257, -0.2125, -0.8795],
                                      lookat=[2.6172, 2.0475, 1.532],
                                      up=[-0.0694, -0.9768, 0.2024])