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


def cut_image(image, num_cuts_h, num_cuts_v):
    width, height = image.size
    cut_width = width // num_cuts_h
    cut_height = height // num_cuts_v
    parts = []
    for i in range(num_cuts_v):
        for j in range(num_cuts_h):
            left = j * cut_width
            upper = i * cut_height
            right = left + cut_width
            lower = upper + cut_height
            parts.append(image.crop((left, upper, right, lower)))

    return parts

def zip_images(images):
    # Create a BytesIO object to hold the zip file
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zipf:
        for idx, img in enumerate(images):
            # Save each image to the zip file
            img_buffer = BytesIO()
            img = Image.fromarray(img)  # Convert NumPy array to PIL Image
            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 process_image(image, num_cuts_h, num_cuts_v):
    # Split the image into a grid of frames
    cut_parts = cut_image(image, num_cuts_h, num_cuts_v)
    #frames = split_image_grid(image, grid_cols_input, grid_rows_input)
    # Zip the frames into a single zip file
    zip_file = zip_images(cut_parts)
    return zip_file

with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(label="Input Image", type="filepath")
        num_cuts_h = gr.Slider(minimum=2, maximum=8, value=3, step=1, label="Number of horizontal cuts/slices")
        num_cuts_v = gr.Slider(minimum=2, maximum=8, value=3, step=1, label="Number of vertical cuts/slices")
        gif_duration = gr.Slider(minimum=5, maximum=1000, value=150, step=5, label="GIF duration (ms)")
        ping_pong_checkbox = gr.Checkbox(label="Ping-pong animation", value=True)
    zip_output = gr.File(label="Output Zip File")
    process_button = gr.Button("Process Image")

    # process_button.click(process_image, inputs=[image_input, num_cuts_h, num_cuts_v, gif_duration, ping_pong_checkbox], outputs=zip_output)
    process_button.click(process_image, inputs=[image_input, num_cuts_h, num_cuts_v], outputs=zip_output)

demo.launch(show_error=True)