import os
import cv2
import numpy as np
import open3d as o3d
import pyrealsense2 as rs

def depth_to_pointcloud(depth_image, intrinsic):
    # Create Open3D Image from depth map
    o3d_depth = o3d.geometry.Image(depth_image)

    # Get intrinsic parameters
    fx, fy, cx, cy = intrinsic.fx, intrinsic.fy, intrinsic.ppx, intrinsic.ppy

    # Create Open3D PinholeCameraIntrinsic object
    o3d_intrinsic = o3d.camera.PinholeCameraIntrinsic(width=depth_image.shape[1], height=depth_image.shape[0], fx=fx, fy=fy, cx=cx, cy=cy)

    # Create Open3D PointCloud object from depth image and intrinsic parameters
    pcd = o3d.geometry.PointCloud.create_from_depth_image(o3d_depth, o3d_intrinsic)

    return pcd

def rgbd_to_pointcloud(rgb_image, depth_image, intrinsic):
    # Create Open3D Image from depth map
    o3d_rgb = o3d.geometry.Image(rgb_image)
    o3d_depth = o3d.geometry.Image(depth_image)
    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
        o3d_rgb, o3d_depth)

    # Get intrinsic parameters
    fx, fy, cx, cy = intrinsic.fx, intrinsic.fy, intrinsic.ppx, intrinsic.ppy

    # Create Open3D PinholeCameraIntrinsic object
    o3d_intrinsic = o3d.camera.PinholeCameraIntrinsic(width=depth_image.shape[1], height=depth_image.shape[0], fx=fx, fy=fy, cx=cx, cy=cy)

    # Create Open3D PointCloud object from depth image and intrinsic parameters
    # pcd = o3d.geometry.PointCloud.create_from_depth_image(o3d_depth, o3d_intrinsic)
    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, o3d_intrinsic)
    
    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
    o3d.visualization.draw_geometries([pcd])

    return pcd

def save_pointcloud(pcd, file_name):
    o3d.io.write_point_cloud(file_name, pcd)

def main():
    # 初始化 RealSense 相机
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 15)
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 15)
    profile = pipeline.start(config)
    
    # Declare pointcloud object, for calculating pointclouds and texture mappings
    # pc = rs.pointcloud()

    align_to = rs.stream.color
    align = rs.align(align_to)

    cv2.namedWindow('RealSense', cv2.WINDOW_AUTOSIZE)
    depth_profile = rs.video_stream_profile(profile.get_stream(rs.stream.depth))
    depth_intrinsics = depth_profile.get_intrinsics()

    if not os.path.exists('data'):
        os.makedirs('data')

    subfolders = ['images', 'depths', 'point_clouds']
    for folder in subfolders:
        folder_path = os.path.join('data', folder)
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

    counter = 1

    try:
        while True:
            frames = pipeline.wait_for_frames()
            # 将深度图和rgb图对齐
            aligned_frames = align.process(frames)
            # 获取深度帧
            aligned_depth_frame = aligned_frames.get_depth_frame()
            if not aligned_depth_frame:
                continue
            # 获得rgb帧
            color_frame = aligned_frames.get_color_frame()
            if not aligned_depth_frame:
                continue
            
            # #  Generate the pointcloud and texture mappings
            # points = pc.calculate(aligned_depth_frame)
            # print(points)


            # 转换为numpy数组 shape (H=720, W=1280)
            depth_image = np.asanyarray(aligned_depth_frame.get_data())
            # print(depth_image.shape)
            color_image = np.asanyarray(color_frame.get_data())

            # 获取内参
            color_intrinsics = color_frame.profile.as_video_stream_profile().intrinsics
            # print("color_intrinsicscolor_image: {}".format(color_intrinsics))
            depth_intrinsics  = aligned_depth_frame.profile.as_video_stream_profile().intrinsics
            # print("depth_intrinsics: {}".format(depth_intrinsics))
            # 使用open3d将深度图转换为点云图
            # pc = depth_to_pointcloud(depth_image, depth_intrinsics)
            # pc = rgbd_to_pointcloud(color_image, depth_image, depth_intrinsics)

            cv2.imshow('RealSense', color_image)

            depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.008), cv2.COLORMAP_JET)
            cv2.imshow('depth_color', depth_colormap)

            # 检查是否按下了 't' 键，如果按下了，就保存当前帧的 RGB、深度图和点云
            key = cv2.waitKey(1)
            if key == ord('t'):
                # 保存 RGB 图像
                rgb_file_path = os.path.join('data', 'images', 'rgb_{:04d}.png'.format(counter))
                cv2.imwrite(rgb_file_path, color_image)
                print('color saved', rgb_file_path)

                # 保存深度图像
                depth_file_path = os.path.join('data', 'depths', 'depth_{:04d}.png'.format(counter))
                cv2.imwrite(depth_file_path, depth_image)
                print('depth saved', depth_file_path)

                # 将点云保存为 pcd 文件
                # pcd_file_path = os.path.join('data', 'point_clouds', 'pointcloud_{:04d}.ply'.format(counter))
                # save_pointcloud(pc, pcd_file_path)
                # print('pc saved', pcd_file_path)
                # points.export_to_ply(pcd_file_path, color_frame)

                # 更新计数器
                counter += 1

            # 检查是否按下了 ESC 键，如果按下了，就退出循环
            if key == 27:
                cv2.destroyAllWindows()
                break
    finally:
        pipeline.stop()

if __name__ == '__main__':
    main()