import cv2

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

# 指定文件夹路径
folder_path = 'datas'

# 获取文件夹中的所有文件
files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]

# 根据文件的修改时间排序
sorted_files = sorted(files, key=lambda x: os.path.getmtime(os.path.join(folder_path, x)))


# 创建一个自定义的 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.figure(figsize=(7, 7))
# 打印排序后的文件名
for file in sorted_files:
    data = np.load(os.path.join(folder_path,file))
    depth = data['arr1']
    points = data['arr2']
    color_image = data['arr3']

    h,w,_ = color_image.shape
    points = points.reshape((h,w,3))

    depth[depth > 1300] = 0
    depth[depth < 900] = 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)

    # 绘制背景
    plt.imshow(background, origin='lower', cmap='gray', alpha=0.3)
    # 绘制白色栅格线，并将其置于底层
    plt.grid(True, which='both', axis='both', color='white', linestyle='-', linewidth=1, zorder=0)
    if False:
        plt.subplot(1,2,1)
        plt.imshow(depth, cmap='plasma', vmin=1000, vmax=1200)
        plt.subplot(1,2,2)
    # 绘制等高线图并设置原点在左下角
    # 通过设置 zorder 来控制它们的层级。例如，设置 zorder=2 或更大的值来确保它们位于栅格线之上。
    plt.contourf(depth, levels=200, cmap=mcmap,vmin=900, vmax=1300,origin='upper',zorder=2)
    plt.pause(0.1)  # 暂停0.1秒
    plt.draw()      # 重绘图像
    plt.clf()       # 清除当前图像
    #plt.show()