import cv2
import numpy as np
import gradio as gr
from PIL import Image


def resize_to_512(img: Image.Image) -> Image.Image:
    if img.size != (512, 512):
        return img.resize((512, 512))
    return img


def gaussian_blur(img: Image.Image, kernel_size: int):
    img = resize_to_512(img)
    img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    blurred = cv2.GaussianBlur(img_cv, (kernel_size | 1, kernel_size | 1), 0)
    return cv2.cvtColor(blurred, cv2.COLOR_BGR2RGB)


def lens_blur(img: Image.Image, max_blur_radius: int):
    img = resize_to_512(img)
    original = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    original_rgb = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)

    # Create synthetic depth map
    depth_norm = np.zeros((original.shape[0], original.shape[1]), dtype=np.float32)
    cv2.circle(depth_norm, (original.shape[1] // 2, original.shape[0] // 2), 100, 1, -1)
    depth_norm = cv2.GaussianBlur(depth_norm, (21, 21), 0)

    blurred_image = np.zeros_like(original_rgb)

    for i in range(original.shape[0]):
        for j in range(original.shape[1]):
            blur_radius = int(depth_norm[i, j] * max_blur_radius)
            if blur_radius % 2 == 0:
                blur_radius += 1

            x_min = max(j - blur_radius, 0)
            x_max = min(j + blur_radius, original.shape[1])
            y_min = max(i - blur_radius, 0)
            y_max = min(i + blur_radius, original.shape[0])

            roi = original_rgb[y_min:y_max, x_min:x_max]

            if blur_radius > 1:
                blurred_roi = cv2.GaussianBlur(roi, (blur_radius, blur_radius), 0)
                try:
                    blurred_image[i, j] = blurred_roi[
                        blur_radius // 2, blur_radius // 2
                    ]
                except:
                    blurred_image[i, j] = original_rgb[i, j]
            else:
                blurred_image[i, j] = original_rgb[i, j]

    return blurred_image


with gr.Blocks() as demo:
    gr.Markdown("## Gaussian and Lens Blur App")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload an Image")
    with gr.Row():
        kernel_slider = gr.Slider(1, 49, value=11, step=2, label="Gaussian Kernel Size")
        max_blur_slider = gr.Slider(
            1, 50, value=15, step=1, label="Max Lens Blur Radius"
        )
    with gr.Row():
        gaussian_output = gr.Image(label="Gaussian Blurred Image")
        lens_output = gr.Image(label="Depth-Based Lens Blurred Image")
    with gr.Row():
        blur_btn = gr.Button("Apply Blur")

    blur_btn.click(
        fn=gaussian_blur, inputs=[image_input, kernel_slider], outputs=gaussian_output
    )
    blur_btn.click(
        fn=lens_blur, inputs=[image_input, max_blur_slider], outputs=lens_output
    )

demo.launch()