import gradio as gr
import numpy as np
import cv2

filters = {
    "Sepia": np.array([[0.393, 0.769, 0.189],
                       [0.349, 0.686, 0.168],
                       [0.272, 0.534, 0.131]]),
    "Cool Blue": np.array([[0.272, 0.534, 0.131],
                           [0.349, 0.686, 0.168],
                           [0.393, 0.769, 0.889]]),
    "Warm Orange": np.array([[1.0, 0.5, 0.0],
                             [0.5, 0.8, 0.2],
                             [0.2, 0.3, 0.5]]),
    "Vintage": np.array([[0.6, 0.4, 0.2],
                         [0.3, 0.7, 0.3],
                         [0.2, 0.3, 0.5]]),
    "Vintage Film": np.array([[0.7, 0.5, 0.3],
                              [0.4, 0.6, 0.2],
                              [0.2, 0.3, 0.5]]),
    "Moonlight": np.array([[0.2, 0.4, 0.6],
                           [0.3, 0.5, 0.7],
                           [0.4, 0.6, 0.8]]),
    "Desert Heat": np.array([[0.9, 0.7, 0.3],
                             [0.7, 0.5, 0.2],
                             [0.5, 0.3, 0.1]]),
    "Purple Dream": np.array([[0.5, 0.3, 0.5],
                              [0.2, 0.4, 0.6],
                              [0.4, 0.2, 0.8]]),
    "Forest Green": np.array([[0.2, 0.6, 0.1],
                              [0.3, 0.7, 0.2],
                              [0.1, 0.5, 0.3]]),
    "Cyberpunk": np.array([[0.8, 0.3, 0.9],
                           [0.2, 0.7, 0.8],
                           [0.9, 0.2, 0.7]]),
    "Pastel Dream": np.array([[0.8, 0.7, 0.9],
                              [0.7, 0.9, 0.8],
                              [0.9, 0.8, 0.7]]),
    "Dramatic Contrast": np.array([[1.2, 0.0, 0.0],
                                   [0.0, 1.2, 0.0],
                                   [0.0, 0.0, 1.2]]),
    "Faded Polaroid": np.array([[0.7, 0.6, 0.5],
                                [0.5, 0.6, 0.5],
                                [0.4, 0.5, 0.6]]),
    "Neon Nights": np.array([[0.7, 0.2, 0.9],
                             [0.3, 0.8, 0.2],
                             [0.9, 0.3, 0.7]]),
    "Underwater": np.array([[0.2, 0.4, 0.6],
                            [0.3, 0.5, 0.7],
                            [0.4, 0.6, 0.8]])
}

def apply_filter(image, filter_type):
    if image is None:
        return None

    if filter_type in filters:
        filtered_img = image.dot(filters[filter_type].T)
        np.clip(filtered_img, 0, 255, out=filtered_img)
        return filtered_img.astype(np.uint8)
    elif filter_type == "High Contrast":
        return cv2.convertScaleAbs(image, alpha=1.5, beta=0)
    elif filter_type == "Grayscale":
        return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    elif filter_type == "Invert":
        return 255 - image
    elif filter_type == "Vignette":
        rows, cols = image.shape[:2]
        kernel_x = cv2.getGaussianKernel(cols, cols/4)
        kernel_y = cv2.getGaussianKernel(rows, rows/4)
        kernel = kernel_y * kernel_x.T
        mask = 255 * kernel / np.linalg.norm(kernel)
        vignette = np.copy(image)
        for i in range(3):
            vignette[:,:,i] = vignette[:,:,i] * mask
        return vignette.astype(np.uint8)
    else:  # Original
        return image

filter_options = ["Original"] + list(filters.keys()) + ["High Contrast", "Grayscale", "Invert", "Vignette"]

def update_filter_samples(image):
    if image is None:
        return None, gr.Gallery(), gr.Radio(choices=filter_options, value="Original")
    
    filtered_images = [apply_filter(image, f) for f in filter_options]
    return (filtered_images[0], 
            gr.Gallery(value=[(img, name) for img, name in zip(filtered_images, filter_options)]),
            gr.Radio(choices=filter_options, value="Original"))

def update_output(image, filter_type):
    if image is None or filter_type is None:
        return None
    return apply_filter(image, filter_type)

def gallery_select(evt: gr.SelectData, filter_choice):
    if evt is None:
        return filter_choice
    return filter_options[evt.index]

with gr.Blocks() as demo:
    gr.Markdown("# Color Filter Application")
    gr.Markdown("Upload an image and apply various color filters to it.")
    
    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Upload an image")
            output_image = gr.Image(label="Filtered Image")
        
        with gr.Column(scale=1):
            filter_gallery = gr.Gallery(label="Filter Samples")
            filter_choice = gr.Radio(
                choices=filter_options,
                value="Original",
                label="Select a filter",
                interactive=True
            )
    
    input_image.change(
        fn=update_filter_samples,
        inputs=[input_image],
        outputs=[output_image, filter_gallery, filter_choice]
    )
    
    filter_gallery.select(
        fn=gallery_select,
        inputs=[filter_choice],
        outputs=filter_choice
    )
    
    filter_choice.change(
        fn=update_output,
        inputs=[input_image, filter_choice],
        outputs=output_image
    )

demo.launch()