import gradio as gr
from PIL import Image, ImageDraw, ImageFont, ImageSequence
import numpy as np

def create_grid(image, grid_size=3, frame_count=9):
    img = Image.open(image.name)

    # Get the total number of frames in the gif
    total_frames = img.n_frames

    # Calculate the frames to use for the grid
    selected_frames_indices = np.linspace(0, total_frames - 1, frame_count, dtype=int)

    # Create a blank image to hold the grid
    grid_dim = (img.width * grid_size, img.height * grid_size)
    grid_img = Image.new("RGBA", grid_dim, "white")

    # Font setup for numbering frames
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 40)
    except IOError:
        font = ImageFont.load_default()

    draw = ImageDraw.Draw(grid_img)

    # Extract frames and place them in the grid
    for i, frame_index in enumerate(selected_frames_indices):
        img.seek(frame_index)
        frame = img.copy()
        x_offset = (i % grid_size) * img.width
        y_offset = (i // grid_size) * img.height
        grid_img.paste(frame, (x_offset, y_offset))

        # Draw frame number
        text = str(i + 1)
        text_position = (x_offset + 20, y_offset + 20)
        text_color = "black" if np.mean(np.array(frame)[:,:3]) > 128 else "white"
        draw.text(text_position, text, fill=text_color, font=font)

    return grid_img

iface = gr.Interface(
    fn=create_grid, 
    inputs=[
        gr.inputs.File(label="Upload a GIF"), 
        gr.inputs.Number(label="Grid Size", default=3, min_value=1, max_value=10),
        gr.inputs.Number(label="Frame Count", default=9, min_value=1, max_value=50)
    ], 
    outputs=gr.outputs.Image(label="Generated Grid"),
    live=False
)

iface.launch()