import mediapy
import numpy as np
from pathlib import Path
import os
from PIL import Image
from tqdm import tqdm

# Create output directory
output_dir = Path("runs/debug_filtered_gifs")
output_dir.mkdir(parents=True, exist_ok=True)

# Find all mp4 files starting with "replay"
input_dir = Path("runs/debug_filtered")
for video_path in tqdm(input_dir.glob("replay*.mp4")):
    # Load video
    video = mediapy.read_video(str(video_path))
    
    # Downsample frames (keep every 6th frame)
    video = video[::6]
    
    # Resize frames to 1/2 size
    h, w = video.shape[1:3]
    new_h, new_w = h//2, w//2
    
    # Convert frames to PIL images and resize
    frames = [Image.fromarray(frame).resize((new_w, new_h)) for frame in video]

    # Save as GIF
    output_path = output_dir / f"{video_path.stem}.gif"
    frames[0].save(
        str(output_path),
        save_all=True,
        append_images=frames[1:],
        duration=1000//5,  # Frame duration in ms (for 5 fps) 
        loop=0
    )