import gradio as gr
import torch
import torch.nn.functional as F
from transformers import CLIPProcessor, CLIPModel

# 1. Load the Engine (Runs once when the server starts)
model_id = "openai/clip-vit-base-patch32"
model = CLIPModel.from_pretrained(model_id)
processor = CLIPProcessor.from_pretrained(model_id)

# 2. Load the Database
# This loads the mathematical vectors you generated in Colab
artist_vectors = torch.load("artist_database.pt")

# 3. The Core Logic Function
def detect_mimicry(suspect_image):
    # Process the uploaded image
    inputs = processor(images=suspect_image, return_tensors="pt")
    with torch.no_grad():
        suspect_vector = model.get_image_features(**inputs)
        
        if not isinstance(suspect_vector, torch.Tensor):
            if hasattr(suspect_vector, "pooler_output"):
                suspect_vector = suspect_vector.pooler_output
            elif hasattr(suspect_vector, "image_embeds"):
                suspect_vector = suspect_vector.image_embeds
                
    # Calculate Cosine Similarity against EVERY image in the database
    similarities = F.cosine_similarity(suspect_vector, artist_vectors)
    
    # Find the highest score (the closest match)
    max_score = torch.max(similarities).item()
    
    # Format the output for the law students
    if max_score > 0.85:
        verdict = f"⚠️ HIGH RISK: Similarity Score of {max_score:.4f}"
    else:
        verdict = f"✅ LOW RISK: Similarity Score of {max_score:.4f}"
        
    return verdict

custom_css = """
[id="gradio-share-link-button"] { display: none !important; }

.share-wrap { display: none !important; }
"""

# 4. Build the Web Interface with Blocks (The Professional Way)
with gr.Blocks(title="⚖️ Artist Copyright Detector") as interface:
    
    # The Header
    gr.Markdown("# ⚖️ Artist Copyright Detector")
    gr.Markdown("Upload an image file to calculate its mathematical similarity against our protected artist database.")
    
    # The Layout (Side-by-Side)
    with gr.Row():
        
        # Left Column: User Input
        with gr.Column():
            image_in = gr.Image(type="pil", sources=["upload"], label="Upload Suspect AI Image")
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Submit", variant="primary")
                
        # Right Column: The Engine Output
        with gr.Column():
            text_out = gr.Text(label="Verdict & Score")
            
    # --- The Wiring ---
    submit_btn.click(fn=detect_mimicry, inputs=image_in, outputs=text_out)
    clear_btn.click(lambda: (None, ""), inputs=None, outputs=[image_in, text_out])

# 5. Launch the app
interface.launch(theme=gr.themes.Monochrome())