import spaces
import numpy as np
from pixeloe.pixelize import pixelize
from PIL import Image
import gradio as gr

def pixelize_image(
    image,
    downscale_mode="contrast",
    target_size=128,
    patch_size=16,
    thickness=2,
    color_matching=True,
    upscale=True
):
    """
    Apply pixelization effect to an image or batch of images.
    """
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    processed = pixelize(
        image,
        mode=downscale_mode,
        target_size=target_size,
        patch_size=patch_size,
        thickness=thickness,
        contrast=1.0,
        saturation=1.0,
        color_matching=color_matching,
        no_upscale=not upscale
    )
    
    return Image.fromarray(processed)

@spaces.GPU
def process_image(image, downscale_mode, target_size, patch_size, thickness, color_matching, upscale):
    if image is None:
        return None
    
    result = pixelize_image(
        image,
        downscale_mode=downscale_mode,
        target_size=target_size,
        patch_size=patch_size,
        thickness=thickness,
        color_matching=color_matching,
        upscale=upscale
    )
    
    return result

def create_pixelize_tab():
    with gr.Tab("Pixelizer"):
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="pil", height=256)
                
                downscale_mode = gr.Dropdown(
                    choices=["contrast", "bicubic", "nearest", "center", "k-centroid"],
                    value="contrast",
                    label="Downscale Mode"
                )
                
                target_size = gr.Slider(
                    minimum=8,
                    maximum=1024,
                    value=128,
                    step=8,
                    label="Target Size"
                )
                
                patch_size = gr.Slider(
                    minimum=4,
                    maximum=32,
                    value=16,
                    step=2,
                    label="Patch Size"
                )
                
                thickness = gr.Slider(
                    minimum=1,
                    maximum=16,
                    value=2,
                    step=1,
                    label="Thickness"
                )
                
                color_matching = gr.Checkbox(
                    value=True,
                    label="Color Matching"
                )
                
                upscale = gr.Checkbox(
                    value=True,
                    label="Upscale"
                )
                
                process_btn = gr.Button("Process Image")
            
            with gr.Column():
                output_image = gr.Image(label="Processed Image", height=256)
        
        # Set up processing event
        process_btn.click(
            fn=process_image,
            inputs=[
                input_image,
                downscale_mode,
                target_size,
                patch_size,
                thickness,
                color_matching,
                upscale
            ],
            outputs=output_image
        )