import gradio as gr
import numpy as np
import zipfile
import io
from PIL import Image

def split_image_grid(image, grid_size):
    # Convert the image to a NumPy array
    img = np.array(image)
    width, height = img.shape[1], img.shape[0]
    grid_width, grid_height = grid_size, grid_size

    # Calculate the size of each grid cell
    cell_width = width // grid_width
    cell_height = height // grid_height

    # Split the image into individual frames
    frames = []
    for i in range(grid_height):
        for j in range(grid_width):
            left = j * cell_width
            upper = i * cell_height
            right = (j + 1) * cell_width
            lower = (i + 1) * cell_height
            frame = img[upper:lower, left:right]
            frames.append(frame)

    return frames

def create_zip_file(frames):
    # Create an in-memory zip file
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zipf:
        for idx, frame in enumerate(frames):
            frame_byte_array = io.BytesIO()
            # Save the frame as a PNG file in the byte array
            frame_img = Image.fromarray(frame)
            frame_img.save(frame_byte_array, format="PNG")
            zipf.writestr(f"frame_{idx}.png", frame_byte_array.getvalue())

    zip_buffer.seek(0)
    return zip_buffer

def create_gif(frames):
    # Create a GIF from the frames
    gif_buffer = io.BytesIO()
    frames_pil = [Image.fromarray(frame) for frame in frames]
    frames_pil[0].save(gif_buffer, format="GIF", save_all=True, append_images=frames_pil[1:], loop=0)
    gif_buffer.seek(0)  # Ensure the buffer is at the beginning
    return gif_buffer

def process_image(image, grid_size):
    # Split the image into frames
    frames = split_image_grid(image, grid_size)
    # Create a zip file from the frames
    zip_file = create_zip_file(frames)
    # Create a GIF from the frames
    gif_file = create_gif(frames)
    return zip_file, gif_file

with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(type="pil", label="Input Image")
        grid_size_slider = gr.Slider(1, 10, value=2, step=1, label="Grid Size")
    with gr.Row():
        zip_output = gr.File(label="Download ZIP")
        gif_output = gr.Image(label="Generated GIF")
    process_button = gr.Button("Process Image")

    process_button.click(process_image, inputs=[image_input, grid_size_slider], outputs=[zip_output, gif_output])

demo.launch(show_error=True)