import gradio as gr
import pandas as pd
from utils.search_functions import run_search

if __name__ == "__main__":
# Define the Gradio interface
    with gr.Blocks() as demo:
        gr.Markdown("### Search Demo with Side-by-Side Comparison")

        with gr.Row():
            with gr.Column():
                search_instructions = """
                **Choose the Type of Search**
                Decide how you'd like to search:
                - **Text**: Type in a description of the fashion item you're looking for.
                - **Image**: Upload an image of a fashion item you want to find similar results for.
                """
                embedding_instructions = """
                **Select How to Match Results**
                Pick the type of comparison you'd like:
                - **Text-based**: Find items based on the descriptions provided to the model.
                - **Image-based**: Find items based on the images provided to the model.
                - **Both (Balanced)**: Search using a combination of both text and image. This balances between the two equally.
                - **Both (Image-focused)**: Emphasize the image more than the text in the search for closer matches.
                """
                gr.Markdown(search_instructions)
            with gr.Column():
                gr.Markdown(embedding_instructions)

        with gr.Row():
            query_type = gr.Radio(["text", "image"], label="Query Type", value="text")
            embedding_type = gr.Radio(["text", "image", "average", "weighted average"], label="Embedding Type")

        # Input fields
        query_input_text = gr.Textbox(label="Text Query", placeholder="Enter text query here", visible=True)
        query_input_image = gr.Image(label="Upload Image", type="filepath", visible=False)
        search_button = gr.Button("Search")

        viewing_instructions = """
        **Viewing Your Results**
        The results are displayed in a scrollable gallery for each model. To explore each image more closely:
        - Scroll down to browse through the images.
        - Click on any image to open it in full view.
        - Once opened, you can flip through the images using the arrows or swipe gestures.
        Each image in the gallery contains the following information:
        - Similarity Score: Shows how closely the item matches your search (higher is better).
        - Metadata: Additional information about each item, such as category, brand, and material, is displayed beneath each image.
        """
        with gr.Row():
            gr.Markdown(viewing_instructions)

        with gr.Row():
            with gr.Column():
                gr.Markdown("#### FashionCLIP Results")
                # Set the gallery to display 4 images per row (scrollable)
                fclip_output_images = gr.Gallery(label="FashionCLIP Images", columns=5)
                fclip_output_results = gr.Dataframe(label="FashionCLIP Search Results")

            with gr.Column():
                gr.Markdown("#### FashionSigLip Results")
                # Set the gallery to display 4 images per row (scrollable)
                siglip_output_images = gr.Gallery(label="FashionSigLIP Images", columns=5)
                siglip_output_results = gr.Dataframe(label="FashionSigLIP Search Results")

        # Update input fields based on query type
        def update_query_input(query_type):
            if query_type == "text":
                return gr.update(visible=True), gr.update(visible=False)
            else:
                return gr.update(visible=False), gr.update(visible=True)

        query_type.change(update_query_input, inputs=query_type, outputs=[query_input_text, query_input_image])

        # Error handling function
        def validate_and_search(query_type, embedding_type, query_input_text, query_input_image):
            if query_type == "text" and not query_input_text:
                return gr.Error("Please enter a text query"), None, None, None
            elif query_type == "image" and not query_input_image:
                return gr.Error("Please upload an image"), None, None, None
            else:
                return run_search(query_type, embedding_type, query_input_text, query_input_image)

        # Search button to trigger both searches with validation
        search_button.click(
            validate_and_search,
            inputs=[query_type, embedding_type, query_input_text, query_input_image],
            outputs=[fclip_output_images, fclip_output_results, siglip_output_images, siglip_output_results]
        )

    # Launch the Gradio app
    demo.launch(share=True)