import gradio as gr
import spaces
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import CLIPProcessor, CLIPModel, SegformerImageProcessor, SegformerForSemanticSegmentation

# ==========================================
# 1. Load Models (CLIP for Vocab, SegFormer for Bias)
# ==========================================
print("Loading models...")
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")

seg_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
seg_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
print("Models loaded successfully.")

# ==========================================
# 2. Core Logic Functions
# ==========================================
def get_clip_saliency(image, text_queries):
    """Extracts per-patch similarity between image and text queries."""
    inputs = clip_processor(text=text_queries, images=image, return_tensors="pt", padding=True)
    
    with torch.no_grad():
        outputs = clip_model(**inputs)
    
    # Get image patch embeddings (Shape: [1, 196, 768] -> 14x14 patches)
    image_embeds = outputs.image_embeds # This is the CLS token (global)
    
    # To get patch-level saliency, we need the last hidden state before projection
    vision_outputs = clip_model.vision_model(pixel_values=inputs['pixel_values'])
    patch_embeds = vision_outputs.last_hidden_state[:, 1:, :] # Exclude CLS token
    patch_embeds = clip_model.visual_projection(patch_embeds) # Project to text space
    
    text_embeds = outputs.text_embeds # Shape: [num_texts, 512]
    
    # Normalize
    patch_embeds = patch_embeds / patch_embeds.norm(dim=-1, keepdim=True)
    text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
    
    # Calculate similarity matrix
    similarity = torch.matmul(patch_embeds[0], text_embeds.T) # Shape: [196, num_texts]
    
    # Reshape to 14x14 spatial dimensions
    saliency_maps = similarity.reshape(14, 14, len(text_queries)).cpu().numpy()
    
    return saliency_maps

def get_foreground_bias(image):
    """Segments the image to find the main object vs background."""
    inputs = seg_processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = seg_model(**inputs)
    
    # Upscale segmentation to original image size
    logits = outputs.logits
    upscaled_logits = torch.nn.functional.interpolate(
        logits, size=(image.size[1], image.size[0]), mode="bilinear", align_corners=False
    )
    
    # Get the most prominent object class (ignoring background class 0 in ADE20k)
    seg_map = upscaled_logits.argmax(dim=1)[0].cpu().numpy()
    
    # Create a binary mask: 1 for foreground objects, 0 for background
    foreground_mask = (seg_map > 0).astype(np.uint8) * 255
    return foreground_mask

@spaces.GPU
def analyze_transparency(image, text_prompts, opacity_slider):
    if not image or not text_prompts:
        return None, None, "Please provide an image and text prompts."
    
    image = image.convert("RGB")
    text_queries = [t.strip() for t in text_prompts.split(",") if t.strip()]
    
    # 1. Get Saliency Maps
    saliency_maps = get_clip_saliency(image, text_queries)
    
    # 2. Get Foreground Bias Mask
    fg_mask = get_foreground_bias(image)
    fg_pixels = (fg_mask > 0).sum()
    bg_pixels = (fg_mask == 0).sum()
    
    results = []
    bias_reports = []
    
    for i, query in enumerate(text_queries):
        # Extract and normalize the saliency map for this query
        saliency = saliency_maps[:, :, i]
        saliency = (saliency - saliency.min()) / (saliency.max() - saliency.min() + 1e-8)
        
        # Resize to original image size
        saliency_resized = cv2.resize(saliency, (image.size[0], image.size[1]))
        
        # Generate Heatmap
        heatmap = cv2.applyColorMap(np.uint8(255 * saliency_resized), cv2.COLORMAP_INFERNO)
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
        
        # Overlay
        alpha = opacity_slider
        overlay = cv2.addWeighted(np.array(image), 1 - alpha, heatmap, alpha, 0)
        overlay_img = Image.fromarray(overlay)
        
        # Calculate Bias
        # How much attention is on the foreground vs background?
        total_attention = saliency_resized.sum() + 1e-8
        fg_attention = (saliency_resized * (fg_mask > 0)).sum() / total_attention * 100
        bg_attention = (saliency_resized * (fg_mask == 0)).sum() / total_attention * 100
        
        # Bias Flag
        bias_flag = "✅ Normal" if fg_attention > 60 else "⚠️ HIGHLY BIASED"
        if fg_attention < 40:
            bias_flag = "🚨 SEVERE CONTEXT BIAS"
            
        results.append((overlay_img, f"{query}"))
        bias_reports.append(f"**{query}**\n- Foreground Focus: `{fg_attention:.1f}%`\n- Background Focus: `{bg_attention:.1f}%`\n- Status: {bias_flag}\n")
    
    # Combine into a gallery and a text report
    bias_report_text = "### 📊 Contextual Bias Report\n" + "\n".join(bias_reports)
    
    return results, bias_report_text, bias_report_text

# ==========================================
# 3. Gradio UI (Handsome & Modern)
# ==========================================
custom_theme = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"]
).set(
    body_background_fill="linear-gradient(to bottom right, #0f172a, #1e293b)",
    body_text_color="#e2e8f0",
    button_primary_background_fill="#10b981",
    button_primary_background_fill_hover="#059669",
    block_border_width="1px",
    block_border_color="#334155",
    block_background_fill="#1e293b"
)

with gr.Blocks(theme=custom_theme, css="footer {visibility: hidden}") as demo:
    gr.Markdown(
        """
        # 🧠 Aura: Open-Vocabulary Neural Transparency
        <p style="color:#94a3b8;">Type any concept. Discover exactly where the AI looks, and detect if it is cheating by relying on background context.</p>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Input Image", elem_id="input-img")
            text_input = gr.Textbox(
                label="Concepts to Search (comma separated)", 
                placeholder="e.g., a dog, a frisbee, green grass, a red car",
                value="a dog, green grass, a person"
            )
            opacity = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Heatmap Opacity")
            analyze_btn = gr.Button("🔍 Analyze Transparency", variant="primary")
            
        with gr.Column(scale=2):
            gallery = gr.Gallery(
                label="Concept Saliency Maps (Where the AI looks for each concept)", 
                show_label=True, 
                columns=2, 
                object_fit="contain",
                height="500px"
            )
            
    with gr.Row():
        bias_report = gr.Markdown(label="Bias Analysis")

    analyze_btn.click(
        fn=analyze_transparency,
        inputs=[input_image, text_input, opacity],
        outputs=[gallery, bias_report]
    )

if __name__ == "__main__":
    demo.launch()