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

def sepia_and_zip(input_img, num_copies):
    # Apply the sepia filter to the image
    sepia_filter = np.array([
        [0.393, 0.769, 0.189],
        [0.349, 0.686, 0.168],
        [0.272, 0.534, 0.131]
    ])
    input_img = input_img / 255.0
    sepia_img = np.dot(input_img[..., :3], sepia_filter.T)
    sepia_img = np.clip(sepia_img, 0, 1) * 255
    sepia_imgs = [(sepia_img).astype(np.uint8) for _ in range(num_copies)]

    # Convert numpy images to PIL images for gallery
    pil_images = [Image.fromarray(img) for img in sepia_imgs]

    # Prepare ZIP file
    zip_bytes_io = io.BytesIO()
    with ZipFile(zip_bytes_io, 'w') as zip_file:
        for i, img in enumerate(pil_images):
            img_byte_arr = io.BytesIO()
            img.save(img_byte_arr, format='PNG')
            img_byte_arr = img_byte_arr.getvalue()
            zip_file.writestr(f"sepia_image_{i}.png", img_byte_arr)
    zip_bytes_io.seek(0)

    return pil_images, zip_bytes_io

with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image()
        num_copies = gr.Number(label="Number of Copies", value=1)
    with gr.Row():
        gallery = gr.Gallery(label="Sepia Images")
        download_btn = gr.File(label="Download ZIP")
        
    btn = gr.Button("Generate Sepia Images")
    btn.click(fn=sepia_and_zip, 
              inputs=[input_img, num_copies], 
              outputs=[gallery, download_btn])

if __name__ == "__main__":
    demo.launch()