import cv2

import pykinect_azure as pykinect
import matplotlib.pyplot as plt
import numpy as np
import pdb

import os
from matplotlib.colors import LinearSegmentedColormap,ListedColormap
from matplotlib.animation import FuncAnimation, FFMpegWriter


# Initialize the library, if the library is not found, add the library path as argument
pykinect.initialize_libraries()

# Modify camera configuration
device_config = pykinect.default_configuration
device_config.color_format = pykinect.K4A_IMAGE_FORMAT_COLOR_BGRA32
device_config.color_resolution = pykinect.K4A_COLOR_RESOLUTION_720P
device_config.depth_mode = pykinect.K4A_DEPTH_MODE_NFOV_2X2BINNED

# Start device
device = pykinect.start_device(config=device_config)

# 创建一个自定义的 colormap
colors = ['red', 'yellow', 'green', 'blue']

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

# 使用交互模式减少闪烁
plt.ion()
fig, ax = plt.subplots(figsize=(7, 7))
cv2.namedWindow('Transformed color',cv2.WINDOW_NORMAL)

framei = 0
while True:
    # Get capture
    capture = device.update()
    
    # Get depth image
    depth_ret, depth = capture.get_depth_image()
    
    # Get the color image in the depth camera axis
    ret_color, color_image = capture.get_color_image()
    
    if not ret_color or not depth_ret:
        continue
        
    h,w,_ = color_image.shape

    depth[depth > 1100] = 0
    depth[depth < 500] = 0
    # depth = depth[50:200,50:210]

    # 背景图
    background = np.ones_like(depth) * 0.5  # 设定灰色背景
    
    # 使用 np.ma.masked_equal() 来屏蔽深度图中的零值。masked_array 中的值不会被绘制，从而避免了零值的显示。
    depth = np.ma.masked_equal(depth, 0)

    # 清除轴内容而不是整个图形
    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)
    
    # 绘制等高线图并设置原点在左下角
    # 通过设置 zorder 来控制它们的层级。例如，设置 zorder=2 或更大的值来确保它们位于栅格线之上。
    ax.contourf(depth, levels=100, cmap=mcmap,vmin=500, vmax=1100,origin='upper',zorder=2)
    
    # 使用更高效的绘图更新方式
    plt.draw()
    plt.pause(0.001)  # 极短暂停时间
    
    # 显示彩色图像
    cv2.imshow('Transformed color', color_image)
    
    # Press q key to stop
    if cv2.waitKey(1) == ord('q'):  
        break
        
cv2.destroyAllWindows()