import gradio as gr
import os
import numpy as np
import torch
import clip
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity

# Load CLIP model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

# Configuration
FLAG_IMAGE_DIR = "./named_flags"  # Update this path
FLAG_EMBEDDINGS_PATH = "./flag_embeddings_1.npy"  # Update this path

# Load precomputed embeddings
flag_embeddings = np.load(FLAG_EMBEDDINGS_PATH, allow_pickle=True).item()

# Get all image paths
image_paths = [
    os.path.join(FLAG_IMAGE_DIR, img)
    for img in os.listdir(FLAG_IMAGE_DIR)
    if img.endswith((".png", ".jpg", ".jpeg"))
]

def get_country_name(image_filename):
    """Extract country name from image filename."""
    return os.path.splitext(os.path.basename(image_filename))[0].upper()

def get_image_embedding(image_path):
    """Get embedding for an input image."""
    image = Image.open(image_path).convert("RGB")
    image_input = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        embedding = model.encode_image(image_input)
    return embedding.cpu().numpy()

def find_similar_flags(image_path, top_n=10):
    """Find similar flags based on cosine similarity."""
    query_embedding = get_image_embedding(image_path)
    similarities = {}

    for flag, embedding in flag_embeddings.items():
        similarity = cosine_similarity(query_embedding, embedding)[0][0]
        similarities[flag] = similarity

    sorted_flags = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
    # Verify if the file exists before adding it to the result
    results = []
    for flag_file, similarity in sorted_flags[1:top_n + 1]:  # Skip the first one as it's the same flag
        flag_path = os.path.join(FLAG_IMAGE_DIR, flag_file)
        if os.path.exists(flag_path):
            results.append((flag_file, similarity))
        else:
            print(f"File not found: {flag_file}")
    return results

def search_flags(query):
    """Search flags based on country name."""
    if not query:
        return image_paths
    return [img for img in image_paths if query.lower() in get_country_name(img).lower()]

def analyze_and_display(selected_flag):
    """Main function to analyze flag similarity and prepare display."""
    try:
        if selected_flag is None:
            return None

        similar_flags = find_similar_flags(selected_flag)
        output_images = []

        for flag_file, similarity in similar_flags:
            flag_path = os.path.join(FLAG_IMAGE_DIR, flag_file)
            country_name = get_country_name(flag_file)
            output_images.append((flag_path, f"{country_name} (Similarity: {similarity:.3f})"))

        return output_images
    except Exception as e:
        return gr.Error(f"Error processing image: {str(e)}")

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Flag Similarity Analysis")
    gr.Markdown("Select a flag from the gallery to find similar flags based on visual features.")

    with gr.Row():
        with gr.Column(scale=1):
            # Search and input gallery
            search_box = gr.Textbox(label="Search Flags", placeholder="Enter country name...")
            input_gallery = gr.Gallery(
                label="Available Flags",
                show_label=True,
                elem_id="gallery",
                columns=4,
                height="auto"
            )

        with gr.Column(scale=1):
            # Output gallery
            output_gallery = gr.Gallery(
                label="Similar Flags",
                show_label=True,
                elem_id="output",
                columns=2,
                height="auto"
            )

    # Event handlers
    def update_gallery(query):
        matching_flags = search_flags(query)
        return [(path, get_country_name(path)) for path in matching_flags]

    def on_select(evt: gr.SelectData, gallery):
        """Handle flag selection from gallery"""
        selected_flag_path = gallery[evt.index][0]
        return analyze_and_display(selected_flag_path)

    # Connect event handlers
    search_box.change(
        update_gallery,
        inputs=[search_box],
        outputs=[input_gallery]
    )

    input_gallery.select(
        on_select,
        inputs=[input_gallery],
        outputs=[output_gallery]
    )

    # Initialize gallery with all flags
    def init_gallery():
        return [(path, get_country_name(path)) for path in image_paths]

    demo.load(init_gallery, outputs=[input_gallery])

# Launch the app
demo.launch()