import os
import imageio
from glob import glob
from collections import defaultdict

# Create output directory for movies
output_dir = "movies"
os.makedirs(output_dir, exist_ok=True)

# Get all PNG files in the current directory
png_files = glob("*.png")

# Group by (A, long)
groups = defaultdict(list)

for file in png_files:
    parts = file.split("_")
    if len(parts) != 3:
        continue  # skip unexpected files
    A = parts[0]               # 'b' or 'pb'
    long = parts[1]            # e.g., '000'
    xxx = parts[2][:-4]        # strip '.png'
    try:
        xxx_val = int(xxx)
        groups[(A, long)].append((xxx_val, file))
    except ValueError:
        continue  # skip if xxx isn't an integer

# Create a movie for each group
for (A, long), frame_list in groups.items():
    # Sort frames numerically by xxx
    sorted_files = [f for _, f in sorted(frame_list)]
    output_file = os.path.join(output_dir, f"{A}_{long}.mp4")

    print(f"Writing {output_file} with {len(sorted_files)} frames...")
    with imageio.get_writer(output_file, fps=5) as writer:
        for img_file in sorted_files:
            image = imageio.imread(img_file)
            writer.append_data(image)

print("All movies written to:", output_dir)