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

class NamedBytesIO(BytesIO):
    def __init__(self, *args, name="output.zip", **kwargs):
        super(NamedBytesIO, self).__init__(*args, **kwargs)
        self.name = name

def split_image_grid(image, grid_cols, grid_rows):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    width, height = image.width, image.height
    cell_width = width // grid_cols
    cell_height = height // grid_rows
    frames = []
    for i in range(grid_rows):
        for j in range(grid_cols):
            left = j * cell_width
            upper = i * cell_height
            right = left + cell_width
            lower = upper + cell_height
            frame = image.crop((left, upper, right, lower))
            frames.append(np.array(frame))
    return frames

def zip_images(images):
    zip_buffer = NamedBytesIO(name="output.zip")
    with zipfile.ZipFile(zip_buffer, 'w') as zipf:
        for idx, img in enumerate(images):
            img_buffer = BytesIO()
            img = Image.fromarray(img)
            img.save(img_buffer, format='PNG')
            img_buffer.seek(0)
            zipf.writestr(f'image_{idx}.png', img_buffer.getvalue())
    zip_buffer.seek(0)
    return zip_buffer

def create_gif(images):
    gif_buffer = NamedBytesIO(name="output.gif")
    images_pil = [Image.fromarray(img) for img in images]
    images_pil[0].save(gif_buffer, format='GIF', save_all=True, append_images=images_pil[1:], duration=100, loop=0)
    gif_buffer.seek(0)
    return gif_buffer

def process_image(image, grid_cols_input, grid_rows_input):
    frames = split_image_grid(image, grid_cols_input, grid_rows_input)
    zip_file = zip_images(frames)
    return zip_file

def process_image_to_gif(image, grid_cols_input, grid_rows_input):
    frames = split_image_grid(image, grid_cols_input, grid_rows_input)
    gif_file = create_gif(frames)
    return gif_file

with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(label="Input Image", type="pil")
        grid_cols_input = gr.Slider(1, 10, value=2, step=1, label="Grid Columns")
        grid_rows_input = gr.Slider(1, 10, value=2, step=1, label="Grid Rows")
    with gr.Row():
        zip_button = gr.Button("Create Zip File")
        gif_button = gr.Button("Create GIF")
    with gr.Row():
        zip_output = gr.File(label="Download Zip File")
        gif_output = gr.File(label="Download GIF")
    zip_button.click(process_image, inputs=[image_input, grid_cols_input, grid_rows_input], outputs=zip_output)
    gif_button.click(process_image_to_gif, inputs=[image_input, grid_cols_input, grid_rows_input], outputs=gif_output)

demo.launch(show_error=True)