import numpy as np
import gradio as gr
import torch
from PIL import Image, ImageFilter
from transformers import AutoImageProcessor, SegformerForSemanticSegmentation, pipeline

# Models
SEG_MODEL_NAME = "nvidia/segformer-b0-finetuned-ade-512-512"
DEPTH_MODEL_NAME = "LiheYoung/depth-anything-base-hf"
DEVICE = "cpu"

seg_processor = AutoImageProcessor.from_pretrained(SEG_MODEL_NAME)
seg_model = SegformerForSemanticSegmentation.from_pretrained(SEG_MODEL_NAME).to(DEVICE)
seg_model.eval()

depth_pipe = pipeline(
    task="depth-estimation",
    model=DEPTH_MODEL_NAME,
    device=-1
)

PERSON_CLASS_ID = 12


# Helpers
def get_person_mask(image: Image.Image) -> np.ndarray:
    image = image.convert("RGB")
    w, h = image.size
    inputs = seg_processor(images=image, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        outputs = seg_model(**inputs)

    logits = outputs.logits
    upsampled_logits = torch.nn.functional.interpolate(
        logits,
        size=(h, w),
        mode="bilinear",
        align_corners=False
    )
    pred_seg = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
    binary_mask = (pred_seg == PERSON_CLASS_ID).astype(np.uint8) * 255
    return binary_mask


def gaussian_background_blur(image: Image.Image, sigma: float = 15.0):
    image = image.convert("RGB")
    mask = get_person_mask(image)
    image_np = np.array(image).astype(np.float32)
    mask_np = (mask.astype(np.float32) / 255.0)
    mask_3 = np.stack([mask_np] * 3, axis=-1)
    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=sigma))
    blurred_np = np.array(blurred_image).astype(np.float32)
    output = mask_3 * image_np + (1.0 - mask_3) * blurred_np
    output = np.clip(output, 0, 255).astype(np.uint8)
    return image, Image.fromarray(mask), Image.fromarray(output)


def depth_based_lens_blur(image: Image.Image):
    image = image.convert("RGB").resize((512, 512))

    # segmentation mask
    mask = get_person_mask(image)
    mask_np = mask.astype(np.float32) / 255.0

    # depth map
    result = depth_pipe(image)
    depth_output = result["depth"]
    depth_np = np.array(depth_output).astype(np.float32)
    depth_norm = (depth_np - depth_np.min()) / (depth_np.max() - depth_np.min() + 1e-8)
    depth_norm = 1.0 - depth_norm
    depth_smooth = np.array(
        Image.fromarray((depth_norm * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(radius=2))
    ).astype(np.float32) / 255.0
    depth_final = depth_smooth * (1.0 - 0.85 * mask_np)
    blur_radii = [0, 2, 4, 6, 8, 10]
    blurred_stack = []

    for r in blur_radii:
        blurred = image.filter(ImageFilter.GaussianBlur(radius=r))
        blurred_stack.append(np.array(blurred).astype(np.float32))

    blurred_stack = np.stack(blurred_stack, axis=0)
    num_levels = len(blur_radii)
    blur_pos = depth_final * (num_levels - 1)
    low_idx = np.floor(blur_pos).astype(np.int32)
    high_idx = np.clip(low_idx + 1, 0, num_levels - 1)
    alpha = blur_pos - low_idx
    image_np = np.array(image).astype(np.float32)
    output = np.zeros_like(image_np)

    for y in range(image_np.shape[0]):
        for x in range(image_np.shape[1]):
            li = low_idx[y, x]
            hi = high_idx[y, x]
            a = alpha[y, x]
            output[y, x] = (1 - a) * blurred_stack[li, y, x] + a * blurred_stack[hi, y, x]

    output = np.clip(output, 0, 255).astype(np.uint8)
    depth_vis = (depth_norm * 255).astype(np.uint8)

    return image, Image.fromarray(depth_vis), Image.fromarray(output)


# Gradio wrappers
def run_gaussian(image):
    if image is None:
        return None, None, None
    return gaussian_background_blur(image, sigma=15.0)


def run_depth_blur(image):
    if image is None:
        return None, None, None
    return depth_based_lens_blur(image)


# UI
with gr.Blocks(title="Problem 2 Blur Demo") as demo:
    gr.Markdown("# Problem 2: Gaussian Blur and Depth-Based Lens Blur")
    gr.Markdown("Upload an image with a clear foreground subject. Then run either Gaussian background blur or depth-based lens blur.")

    with gr.Row():
        input_image = gr.Image(type="pil", label="Input Image")

    with gr.Row():
        gaussian_btn = gr.Button("Run Gaussian Background Blur")
        depth_btn = gr.Button("Run Depth-Based Lens Blur")

    gr.Markdown("## Gaussian Background Blur")
    with gr.Row():
        g_input = gr.Image(type="pil", label="Input")
        g_mask = gr.Image(type="pil", label="Foreground Mask")
        g_output = gr.Image(type="pil", label="Blurred Background")

    gr.Markdown("## Depth-Based Lens Blur")
    with gr.Row():
        d_input = gr.Image(type="pil", label="Input (512x512)")
        d_depth = gr.Image(type="pil", label="Depth Map")
        d_output = gr.Image(type="pil", label="Lens Blur Output")

    gaussian_btn.click(
        fn=run_gaussian,
        inputs=input_image,
        outputs=[g_input, g_mask, g_output]
    )

    depth_btn.click(
        fn=run_depth_blur,
        inputs=input_image,
        outputs=[d_input, d_depth, d_output]
    )

demo.launch(share=True)