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)


# Load model once globally
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)

    # Get depth map
    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()
    )

    # Normalize and invert depth
    depth_norm = (depth - depth.min()) / (depth.max() - depth.min())
    depth_inverted = 1.0 - depth_norm

    # Dynamically scale blur strength using the slider
    num_levels = 6  # More levels for smoother transitions
    max_sigma = (
        max_blur_radius / 2.0
    )  # Scale down to reasonable range (e.g. 0–25 → 0–12.5 sigma)
    blur_levels = np.linspace(0, max_sigma, num_levels)
    blurred_images = [gaussian_filter(original, sigma=(s, s, 0)) for s in blur_levels]

    # Blend based on depth
    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)


# Separate update functions
def update_gaussian(img, kernel_size):
    return gaussian_blur(img, kernel_size)


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


def apply_blurs(img, kernel_size, max_blur_radius):
    g_blurred = gaussian_blur(img, kernel_size)
    l_blurred = lens_blur(img, max_blur_radius)
    return g_blurred, l_blurred


with gr.Blocks() as demo:
    gr.Markdown("## 🌀 Apply Gaussian and Depth-Based Lens Blur")

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

    with gr.Row():
        kernel_slider = gr.Slider(1, 49, step=2, value=11, label="Gaussian Kernel Size")
        lens_slider = gr.Slider(
            1, 50, step=1, value=15, label="Max Lens Blur Intensity"
        )

    with gr.Row():
        gaussian_output = gr.Image(label="Gaussian Blurred Image")
        lens_output = gr.Image(label="Depth-Based Lens Blurred Image")

    # Trigger both when image changes
    image_input.change(
        fn=apply_blurs,
        inputs=[image_input, kernel_slider, lens_slider],
        outputs=[gaussian_output, lens_output],
    )

    # Trigger only gaussian blur
    kernel_slider.change(
        fn=update_gaussian,
        inputs=[image_input, kernel_slider],
        outputs=gaussian_output,
    )

    # Trigger only lens blur
    lens_slider.change(
        fn=update_lens,
        inputs=[image_input, lens_slider],
        outputs=lens_output,
    )

demo.launch()