import os
import numpy as np
import cv2
from matplotlib.colors import LinearSegmentedColormap


class FemtoBoltViewer:
    def __init__(self, depth_min=900, depth_max=1300):
        self.depth_min = depth_min
        self.depth_max = depth_max

        # 自定义彩虹色 colormap
        colors = ['fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue',
                  'fuchsia', 'red', 'yellow', 'lime', 'cyan', 'blue']
        self.cmap = LinearSegmentedColormap.from_list("custom_cmap", colors)

        # SDK 设备句柄和配置
        self.device_handle = None
        self.pykinect = None
        self.config = None

        # 缓存背景+网格图像（仅生成一次）
        self.background = None

        # OpenCV 窗口
        cv2.namedWindow("Depth CV")

    def _load_sdk(self):
        """加载并初始化 FemtoBolt SDK"""
        try:
            import pykinect_azure as pykinect
            self.pykinect = pykinect
            base_dir = os.path.dirname(os.path.abspath(__file__))
            dll_path = os.path.join(base_dir, "..", "dll", "femtobolt", "bin", "k4a.dll")
            self.pykinect.initialize_libraries(track_body=False, module_k4a_path=dll_path)
            return True
        except Exception as e:
            print(f"加载 SDK 失败: {e}")
            return False

    def _configure_device(self):
        """配置 FemtoBolt 深度相机"""
        self.config = self.pykinect.default_configuration
        self.config.depth_mode = self.pykinect.K4A_DEPTH_MODE_NFOV_UNBINNED
        self.config.camera_fps = self.pykinect.K4A_FRAMES_PER_SECOND_15
        self.config.synchronized_images_only = False
        self.config.color_resolution = 0
        self.device_handle = self.pykinect.start_device(config=self.config)

    def _get_color_image(self, depth_image):
        """将原始深度图转换为叠加背景网格后的 RGB 彩色图像"""
        h, w = depth_image.shape
        # 第一次调用时生成灰色背景和白色网格
        if self.background is None:
            self.background = np.full((h, w, 3), 128, dtype=np.uint8)  # 灰色 (0.5 -> 128)
            # 绘制网格线
            for x in range(w):
                cv2.line(self.background, (x, 0), (x, h-1), (255, 255, 255), 1)
            for y in range(h):
                cv2.line(self.background, (0, y), (w-1, y), (255, 255, 255), 1)

        # 生成深度掩码，仅保留指定范围内的像素
        mask_valid = (depth_image >= self.depth_min) & (depth_image <= self.depth_max)
        depth_clipped = np.clip(depth_image, self.depth_min, self.depth_max)
        normed = (depth_clipped.astype(np.float32) - self.depth_min) / (self.depth_max - self.depth_min)

        # 反转映射，保证颜色方向与之前一致
        normed = 1.0 - normed

        # 应用自定义 colormap，将深度值映射到 RGB
        rgba = self.cmap(normed)
        rgb = (rgba[..., :3] * 255).astype(np.uint8)

        # 叠加：在背景上覆盖彩色深度图（掩码处不覆盖，保留灰色背景+网格）
        final_img = self.background.copy()
        final_img[mask_valid] = rgb[mask_valid]
        return final_img

    def run(self):
        if not self._load_sdk():
            print("SDK 加载失败，程序退出")
            return

        self._configure_device()
        print("FemtoBolt 深度相机启动成功，按 Ctrl+C 或 ESC 退出")

        try:
            while True:
                capture = self.device_handle.update()
                if capture is None:
                    continue
                ret, depth_image = capture.get_depth_image()
                if not ret or depth_image is None:
                    continue

                # 转换并渲染当前帧
                final_img = self._get_color_image(depth_image)

                # OpenCV 显示
                cv2.imshow("Depth CV", final_img)
                # 按 ESC 键退出
                if cv2.waitKey(1) & 0xFF == 27:
                    break

        except KeyboardInterrupt:
            print("检测到退出信号，结束程序")
        finally:
            if self.device_handle:
                self.device_handle.stop()
                self.device_handle.close()
            cv2.destroyAllWindows()


if __name__ == "__main__":
    viewer = FemtoBoltViewer(depth_min=900, depth_max=1100)
    viewer.run()