import os
import zipfile
from datetime import datetime
import imageio
from PIL import Image

def concatenate_images_horizontally(image_paths, output_path):
    images = [Image.open(img_path) for img_path in image_paths]
    
    # 计算总宽度和最大高度
    total_width = sum(img.width for img in images)
    max_height = max(img.height for img in images)

    # 创建空白画布
    combined_image = Image.new('RGB', (total_width, max_height))

    # 拼接图片
    x_offset = 0
    for img in images:
        combined_image.paste(img, (x_offset, 0))
        x_offset += img.width

    combined_image.save(output_path)
# start_name = input("start_name:")
# end_name = input("end_name:")

start_name = "/data1/zxx/runs/racegame/level1/swift_qwen2vl_config_7B/2024-11-24-19-10-52-452773/screen_2024-11-24-19-11-03-018045.png"
end_name = "/data1/zxx/runs/racegame/level1/swift_qwen2vl_config_7B/2024-11-24-19-10-52-452773/screen_2024-11-24-19-15-26-657684.png"


assert os.path.dirname(start_name) == os.path.dirname(end_name)

folder_path = os.path.dirname(start_name)

print(folder_path)


# 获取文件的时间信息以排序
def extract_timestamp(filename):
    date_str = filename.replace("screen_", "").replace(".png", "")
    return datetime.strptime(date_str, "%Y-%m-%d-%H-%M-%S-%f")

# 获取文件并排序
png_lst = sorted(
    [f for f in os.listdir(folder_path) if f.endswith('.png') and start_name <= os.path.join(folder_path, f) <= end_name],
    key=extract_timestamp
)
png_lst = [os.path.join(folder_path, f) for f in png_lst]


# 压缩文件
zip_filename = "selected_images.zip"
with zipfile.ZipFile(zip_filename, 'w') as zipf:
    for file in png_lst:
        zipf.write(file, arcname=file)



frames = []
for file in png_lst:
    frames.append(imageio.imread(file))
    
gif_filename = "selected_images.gif"

imageio.mimsave("selected_images.gif", frames, 'GIF', duration=0.1, loop=0)


png_filename = "selected_images.png"

concatenate_images_horizontally(png_lst, png_filename)

print(f"已将 {len(png_lst)} 张图片成功压缩到 {zip_filename} 中，生成{gif_filename}, {png_filename}。")