import os
import io
import numpy as np
import cv2
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

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

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

        # Matplotlib 图形初始化
        self.fig, self.ax = plt.subplots(figsize=(6, 6), dpi=75)
        self.ax.axis('off')  # 隐藏坐标轴

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

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

    def _load_sdk(self):
        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", "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):
        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.device_handle = self.pykinect.start_device(config=self.config)

    def _generate_contour_image_plt(self, depth):
        """使用 matplotlib 生成等高线图像（完全采用 display_x.py 的逻辑）"""
        try:
            # 清除之前的绘图
            self.ax.clear()
            self.ax.axis('off')

            # 深度数据过滤
            depth_filtered = depth.copy()
            depth_filtered[depth_filtered > self.depth_range_max] = 0
            depth_filtered[depth_filtered < self.depth_range_min] = 0

            # 背景图
            background = np.ones_like(depth_filtered) * 0.5  # 灰色背景
            self.ax.imshow(background, origin='lower', cmap='gray', alpha=0.3)

            # 屏蔽深度为0
            depth_masked = np.ma.masked_equal(depth_filtered, 0)

            # 绘制白色栅格线（底层）
            self.ax.grid(True, which='both', axis='both', color='white', linestyle='-', linewidth=0.5, zorder=0)
            self.ax.minorticks_on()
            self.ax.grid(True, which='minor', axis='both', color='white', linestyle='-', linewidth=0.3, zorder=0)

            # 绘制等高线图
            self.ax.contourf(depth_masked, levels=100, cmap=self.custom_cmap,
                             vmin=self.depth_range_min, vmax=self.depth_range_max, origin='upper', zorder=2)

            # 保存到 BytesIO 缓冲区
            buf = io.BytesIO()
            self.fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75)
            buf.seek(0)

            # 转为 numpy 数组
            img_array = np.frombuffer(buf.getvalue(), dtype=np.uint8)
            buf.close()

            # 使用 OpenCV 解码 PNG
            img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

            # 裁剪宽度
            if img is not None:
                height, width = img.shape[:2]
                target_width = round(height // 2)
                if width > target_width:
                    left = (width - target_width) // 2
                    right = left + target_width
                    img = img[:, left:right]
                return img
            else:
                print("无法解码matplotlib生成的PNG图像")
                return None

        except Exception as e:
            print(f"生成等高线图像失败: {e}")
            return None

    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._generate_contour_image_plt(depth_image)
                if final_img is not None:
                    cv2.imshow("Depth CV", final_img)
                    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=500, depth_max=700)
    viewer.run()