"""快速预览示例。

同时显示彩色图与深度伪彩图，便于验证相机连接和图像质量。
"""

import time

import cv2
import numpy as np
from pyorbbecsdk import *

from utils import frame_to_bgr_image

ESC_KEY = 27
MIN_DEPTH = 20  # 20mm
MAX_DEPTH = 10000  # 10000mm


def main():
    """启动实时预览窗口。"""
    pipeline = Pipeline()

    pipeline.start()
    print("Pipeline started successfully. Press 'q' or ESC to exit.")

    # 设置窗口大小，左右并排展示彩色图和深度图。
    window_width = 1280
    window_height = 720
    cv2.namedWindow("QuickStart Viewer", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("QuickStart Viewer", window_width, window_height)

    while True:
        try:
            frames = pipeline.wait_for_frames(100)
            if frames is None:
                continue

            # 读取彩色帧并转换为 OpenCV BGR 图像。
            color_frame = frames.get_color_frame()
            if color_frame is None:
                continue
            color_image = frame_to_bgr_image(color_frame)

            # 读取深度帧并校验格式。
            depth_frame = frames.get_depth_frame()
            if depth_frame is None:
                continue
            if depth_frame.get_format() != OBFormat.Y16:
                print("Depth format is not Y16")
                continue

            # 深度数据预处理：重塑、按 scale 转毫米、过滤区间外像素。
            width = depth_frame.get_width()
            height = depth_frame.get_height()
            scale = depth_frame.get_depth_scale()

            depth_data = np.frombuffer(depth_frame.get_data(), dtype=np.uint16).reshape((height, width))
            depth_data = depth_data.astype(np.float32) * scale
            depth_data = np.where((depth_data > MIN_DEPTH) & (depth_data < MAX_DEPTH), depth_data, 0).astype(np.uint16)

            # 深度伪彩可视化。
            depth_image = cv2.normalize(depth_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)

            # 缩放并拼接：左彩色、右深度。
            color_image_resized = cv2.resize(color_image, (window_width // 2, window_height))
            depth_image_resized = cv2.resize(depth_image, (window_width // 2, window_height))
            combined_image = np.hstack((color_image_resized, depth_image_resized))

            cv2.imshow("QuickStart Viewer", combined_image)

            if cv2.waitKey(1) in [ord("q"), ESC_KEY]:
                break
        except KeyboardInterrupt:
            break

    cv2.destroyAllWindows()
    pipeline.stop()
    print("Pipeline stopped and all windows closed.")


if __name__ == "__main__":
    main()