import cv2
import numpy as np
import gradio as gr
from PIL import Image
from scipy.ndimage import gaussian_filter
from transformers import (
    AutoImageProcessor,
    AutoModelForDepthEstimation,
)
import torch


def resize_to_512(img: Image.Image) -> Image.Image:
    return img.resize((512, 512)) if img.size != (512, 512) else 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)


depth_model_id = "depth-anything/Depth-Anything-V2-Small-hf"
processor = AutoImageProcessor.from_pretrained(depth_model_id)
depth_model = AutoModelForDepthEstimation.from_pretrained(depth_model_id)


def lens_blur(img: Image.Image, max_blur_radius: int):
    img = resize_to_512(img)
    original = np.array(img).astype(np.float32)

    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = depth_model(**inputs)
        predicted_depth = outputs.predicted_depth

    depth = (
        torch.nn.functional.interpolate(
            predicted_depth.unsqueeze(1),
            size=(512, 512),
            mode="bicubic",
            align_corners=False,
        )
        .squeeze()
        .cpu()
        .numpy()
    )

    depth_norm = (depth - depth.min()) / (depth.max() - depth.min())
    depth_inverted = 1.0 - depth_norm

    num_levels = 6
    max_sigma = max_blur_radius / 2.0
    blur_levels = np.linspace(0, max_sigma, num_levels)
    blurred_images = [gaussian_filter(original, sigma=(s, s, 0)) for s in blur_levels]

    blurred_final = np.zeros_like(original, dtype=np.float32)
    depth_scaled = depth_inverted * (num_levels - 1)
    depth_int = np.floor(depth_scaled).astype(int)
    depth_frac = depth_scaled - depth_int

    for i in range(num_levels - 1):
        mask = depth_int == i
        alpha = depth_frac[mask]
        for c in range(3):
            blended = (
                blurred_images[i][..., c][mask] * (1 - alpha)
                + blurred_images[i + 1][..., c][mask] * alpha
            )
            blurred_final[..., c][mask] = blended

    return np.clip(blurred_final, 0, 255).astype(np.uint8)


def synthetic_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)

    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


def apply_all_blurs(img, g_kernel, lens_radius, synthetic_radius):
    g = gaussian_blur(img, g_kernel)
    l = lens_blur(img, lens_radius)
    s = synthetic_lens_blur(img, synthetic_radius)
    return g, l, s


def update_gaussian(img, kernel_size):
    return gaussian_blur(img, kernel_size)


def update_lens(img, radius):
    return lens_blur(img, radius)


def update_synthetic(img, radius):
    return synthetic_lens_blur(img, radius)


with gr.Blocks() as demo:
    gr.Markdown(
        "## 🌀 Blur Effects Comparison: Gaussian, Depth-Based, Synthetic (Depth Based Blur works with bottles)"
    )

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")

    with gr.Row():
        g_slider = gr.Slider(1, 49, step=2, value=11, label="Gaussian Kernel Size")
        lens_slider = gr.Slider(
            1,
            50,
            step=1,
            value=15,
            label="Depth-Based Blur Intensity (Works with bottles)",
        )
        synth_slider = gr.Slider(1, 50, step=1, value=25, label="Synthetic Blur Radius")

    with gr.Row():
        g_output = gr.Image(label="Gaussian Blurred Image")
        l_output = gr.Image(label="Depth-Based Lens Blurred Image")
        s_output = gr.Image(label="Synthetic Depth Lens Blurred Image")

    # Initial image upload updates all three
    image_input.change(
        fn=apply_all_blurs,
        inputs=[image_input, g_slider, lens_slider, synth_slider],
        outputs=[g_output, l_output, s_output],
    )

    # Individual updates for each slider
    g_slider.change(
        fn=update_gaussian, inputs=[image_input, g_slider], outputs=g_output
    )
    lens_slider.change(
        fn=update_lens, inputs=[image_input, lens_slider], outputs=l_output
    )
    synth_slider.change(
        fn=update_synthetic, inputs=[image_input, synth_slider], outputs=s_output
    )

demo.launch()