import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import pykinect_azure as pykinect


class FemtoBoltContourViewer:
    def __init__(self, depth_min=900, depth_max=1100):
        self.depth_min = depth_min
        self.depth_max = depth_max

        # 自定义离散彩虹色，层次明显
        colors = [
            'darkblue', 'blue', 'cyan', 'lime', 'yellow',
            'orange', 'red', 'darkred'
        ]
        self.cmap = ListedColormap(colors)

        self.device_handle = None
        self.pykinect = None
        self.config = None

    def _load_sdk(self):
        """加载并初始化 FemtoBolt SDK"""
        base_dir = os.path.dirname(os.path.abspath(__file__))
        dll_path = os.path.join(base_dir, "..", "dll", "femtobolt", "bin", "k4a.dll")
        self.pykinect = pykinect
        self.pykinect.initialize_libraries(track_body=False, module_k4a_path=dll_path)

    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 run(self):
        self._load_sdk()
        self._configure_device()

        plt.ion()  # 打开交互模式
        fig, ax = plt.subplots(figsize=(7, 7))
        print("FemtoBolt 深度相机启动成功，关闭窗口或 Ctrl+C 退出")

        # 设置离散等高线层次
        levels = np.linspace(self.depth_min, self.depth_max, len(self.cmap.colors) + 1)

        try:
            while plt.fignum_exists(fig.number):  # 窗口存在才继续
                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

                depth = depth_image.astype(np.uint16)
                # 限制深度范围
                depth[depth > self.depth_max] = 0
                depth[depth < self.depth_min] = 0
                depth = depth[0:350, 0:350]

                # 屏蔽无效值
                depth_masked = np.ma.masked_equal(depth, 0)

                # 背景灰色
                background = np.ones_like(depth) * 0.5

                # 绘制
                ax.clear()
                ax.imshow(background, origin='lower', cmap='gray', alpha=0.3)
                ax.grid(True, which='both', axis='both', color='white',
                        linestyle='--', linewidth=1, zorder=0)
                ax.contourf(depth_masked, levels=levels, cmap=self.cmap,
                            vmin=self.depth_min, vmax=self.depth_max,
                            origin='upper', zorder=2)

                plt.pause(0.05)

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


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