import os
import io
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# 数据文件夹
files_dir = "./data_files"

# 获取文件夹下所有的符合条件的数据文件
all_files = os.listdir(files_dir)
log_files = []
for filename in all_files:
    if filename.startswith("runtime_data") and filename.endswith(".log"):
        log_files.append(filename)

# 用于存放Gif图片的每一帧图像
gif_frames = []

# 准备图像和坐标系统
figure, axis = plt.subplots()

# 对每一个文件画图
for file in log_files:
    # 获取数据
    file_data = np.loadtxt(files_dir + "/" + file)
    step = file_data[:, 1]
    time_barrier = file_data[:, 7]
    time_bcast = file_data[:, 8]
    # 重新画图
    plt.cla()
    axis.plot(step, time_bcast, 'o-', label='Bcast')
    axis.plot(step, time_barrier, 'o-', label='Barrier')
    # 设置坐标名称
    axis.set_xlabel("Step")
    axis.set_ylabel("Time")
    # 设置标题、网格线、图注
    axis.grid(True)
    axis.legend()
    axis.set_title("Plot file: " + file)
    # 显示曲线
    plt.draw()

    # 曲线图变成Image对象
    image_buffer = io.BytesIO()
    figure.savefig(image_buffer)
    image_buffer.seek(0)
    image = Image.open(image_buffer)
    gif_frames.append(image)
    # 暂停0.5s
    plt.pause(0.5)

# 保存到gif
image0 = gif_frames[0]
image0.save(files_dir + "/runtime_data_time.gif",
            save_all=True,
            append_images=gif_frames[1:],
            duration=100, loop=0)

# 如果想保留最后一个图的窗口
plt.show()