import gradio as gr
import torch
import numpy as np
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from PIL import Image, ImageFilter
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# ---------------------------
# Depth Estimation Utilities
# ---------------------------
def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
    """
    Loads the LiheYoung/depth-anything-large-hf model and computes a depth map.
    The depth map is normalized, inverted (so that near=0 and far=1),
    and multiplied by the given scale_factor.
    """
    # Load model and processor from pretrained weights
    image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf")
    model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf")
    
    # Prepare image for the model
    inputs = image_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_depth = outputs.predicted_depth

    # Interpolate predicted depth map to match image size
    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=image.size[::-1],  # PIL image size is (width, height)
        mode="bicubic",
        align_corners=False,
    )
    # Normalize for visualization
    depth_min = prediction.min()
    depth_max = prediction.max()
    depth_vis = (prediction - depth_min) / (depth_max - depth_min + 1e-8)
    depth_map = depth_vis.squeeze().cpu().numpy()
    # Invert so that near=0 and far=1, then scale
    depth_map_inverted = 1.0 - depth_map
    depth_map_inverted *= scale_factor
    return depth_map_inverted

# ---------------------------
# Depth-Based Blur Functions
# ---------------------------
def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
    """
    Creates multiple blurred versions of the image (using Gaussian blur with radii from 0 to max_blur)
    and composites them using masks generated from bins of the normalized depth map.
    """
    blur_radii = np.linspace(0, max_blur, num_layers)
    blur_versions = [image.filter(ImageFilter.GaussianBlur(radius)) for radius in blur_radii]
    # Use a fixed range (0 to 1) since the depth map is normalized
    thresholds = np.linspace(0, 1, num_layers + 1)
    final_image = blur_versions[-1]
    for i in range(num_layers - 1, -1, -1):
        mask_array = np.logical_and(
            depth_map >= thresholds[i],
            depth_map < thresholds[i + 1]
        ).astype(np.uint8) * 255
        mask_image = Image.fromarray(mask_array, mode="L")
        final_image = Image.composite(blur_versions[i], final_image, mask_image)
    return final_image

def process_depth_blur(uploaded_image, max_blur_value, scale_factor, num_layers):
    """
    Resizes the uploaded image to 512x512, computes its depth map,
    and applies layered blur based on the depth.
    """
    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.open(uploaded_image)
    image = uploaded_image.convert("RGB").resize((512, 512))
    depth_map = compute_depth_map(image, scale_factor)
    final_image = layered_blur(image, depth_map, int(num_layers), max_blur_value)
    return final_image

# ---------------------------
# Depth Heatmap Functions
# ---------------------------
def create_heatmap(depth_map: np.ndarray, intensity: float) -> Image.Image:
    """
    Applies a colormap to the normalized depth map.
    The 'intensity' slider multiplies the normalized depth values (clipped to [0,1])
    before applying the "inferno" colormap.
    """
    # Multiply depth map by intensity and clip to 0-1 range
    normalized = np.clip(depth_map * intensity, 0, 1)
    colormap = cm.get_cmap("inferno")
    colored = colormap(normalized)  # Returns an RGBA image in [0, 1]
    heatmap = (colored[:, :, :3] * 255).astype(np.uint8)  # drop alpha and convert to [0,255]
    return Image.fromarray(heatmap)

def process_depth_heatmap(uploaded_image, intensity):
    """
    Resizes the uploaded image to 512x512, computes its depth map (with scale factor 1.0),
    and returns a heatmap visualization.
    """
    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.open(uploaded_image)
    image = uploaded_image.convert("RGB").resize((512, 512))
    depth_map = compute_depth_map(image, scale_factor=1.0)
    heatmap_img = create_heatmap(depth_map, intensity)
    return heatmap_img

# --- Segmentation-Based Blur using BEN2 ---
def load_segmentation_model():
    """
    Loads and caches the segmentation model from BEN2.
    Ensure you have ben2 installed and accessible in your path.
    """
    global seg_model, seg_device
    if "seg_model" not in globals():
        from ben2 import BEN_Base  # Import BEN2
        seg_model = BEN_Base.from_pretrained("PramaLLC/BEN2")
        seg_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        seg_model.to(seg_device).eval()
    return seg_model, seg_device

def process_segmentation_blur(uploaded_image, seg_blur_radius: float):
    """
    Processes the image with segmentation-based blur.
    The image is resized to 512x512. A Gaussian blur with the specified radius is applied,
    then the segmentation mask is computed to composite the sharp foreground over the blurred background.
    """
    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.open(uploaded_image)
    image = uploaded_image.convert("RGB").resize((512, 512))
    seg_model, seg_device = load_segmentation_model()
    blurred_image = image.filter(ImageFilter.GaussianBlur(seg_blur_radius))
    
    # Generate segmentation mask (foreground)
    foreground = seg_model.inference(image, refine_foreground=False)
    foreground_rgba = foreground.convert("RGBA")
    _, _, _, alpha = foreground_rgba.split()
    binary_mask = alpha.point(lambda x: 255 if x > 128 else 0, mode="L")
    final_image = Image.composite(image, blurred_image, binary_mask)
    return final_image

# --- Merged Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Depth-Based vs Segmentation-Based Blur")
    with gr.Tabs():
        with gr.Tab("Depth Blur"):
            img_input = gr.Image(type="pil", label="Upload Image")
            blur_slider = gr.Slider(1, 50, value=20, label="Maximum Blur Radius")
            scale_slider = gr.Slider(0.1, 2.0, value=1.0, label="Depth Scale Factor")
            layers_slider = gr.Slider(2, 10, value=5, label="Number of Layers")
            blur_output = gr.Image(label="Depth Blur Result")
            blur_button = gr.Button("Process Depth Blur")
            blur_button.click(
                process_depth_blur, 
                inputs=[img_input, blur_slider, scale_slider, layers_slider], 
                outputs=blur_output
            )
        with gr.Tab("Depth Heatmap"):
            img_input2 = gr.Image(type="pil", label="Upload Image")
            intensity_slider = gr.Slider(0.5, 5.0, value=1.0, label="Heatmap Intensity")
            heatmap_output = gr.Image(label="Depth Heatmap")
            heatmap_button = gr.Button("Generate Depth Heatmap")
            heatmap_button.click(
                process_depth_heatmap, 
                inputs=[img_input2, intensity_slider], 
                outputs=heatmap_output
            )
        with gr.Tab("Segmentation-Based Blur (BEN2)"):
            seg_img = gr.Image(type="pil", label="Upload Image")
            seg_blur = gr.Slider(5, 30, value=15, step=1, label="Segmentation Blur Radius")
            seg_out = gr.Image(label="Segmentation-Based Blurred Image")
            seg_button = gr.Button("Process Segmentation Blur")
            seg_button.click(process_segmentation_blur, inputs=[seg_img, seg_blur], outputs=seg_out)

if __name__ == "__main__":
    # Optionally, set share=True to generate a public link.
    demo.launch(share=True)