import gradio as gr
import pandas as pd
from langchain_community.vectorstores import SKLearnVectorStore
from langchain_community.embeddings import HuggingFaceBgeEmbeddings

cols = [
    "Name",
    "Description",
    "Category",
    "AI-Driven",
    "Champion",
    "Match score",
]

persist_path = "aecaihub.parquet"

model_name = "BAAI/bge-small-en-v1.5"
encode_kwargs = {'normalize_embeddings': True,"show_progress_bar":False,"batch_size":1} # set True to compute cosine similarity
embeddings_function = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    encode_kwargs=encode_kwargs,
    query_instruction="Represent this sentence for searching relevant passages: "
)


vector_store = SKLearnVectorStore(
    embedding=embeddings_function, persist_path=persist_path, serializer="parquet"
)

def predict(query,k):

    docs = vector_store.similarity_search_with_score(query,k = k)

    df_results = []
    for doc,score in docs:
        m = doc.metadata
        result_doc = {
            "Name":f"**[{m['name']}]({m['url']})**",
            "Description":doc.page_content,
            "Category":m["category"],
            "AI-Driven":m["ai_driven"],
            "Champion":m["champion"],
            "Match score":round(1-score,3),
        }
        
        df_results.append(result_doc)
    df_results = pd.DataFrame(df_results)
    return df_results


examples = [
    "Tool to generate floor plans"
    "AI tool for comparing building materials and sustainability",
    "3D model library with image search function",
    "AI-powered 3D design tool for architects and interior designers",
    "Software for extracting 3D models from videos",
    "AI tool for comprehensive utility data in infrastructure projects",
    "AI for generating creative content in design projects",
    "AI tool to convert text into architectural videos",
    "AI solutions for low carbon design and data mining in architecture",
    "Software for construction quantity estimation and progress tracking",
    "AI interior design tool for automatic room designs"
]

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("""
        # 🏢 AEC AI Hub - Semantic Search Engine
        This tool uses semantic search to find AI tools for Architecture Engineering and Construction (AEC) based on your question.
        The database is drawn from the [great work](https://stjepanmikulic.notion.site/AEC-AI-Hub-b6e6eebe88094e0e9b4995da38e96768) of [Stjepan Mikulic](https://www.linkedin.com/in/stjepanmikulic/)
    """)

    with gr.Row():
        search_bar = gr.Textbox(label="Ask you question here",scale = 2)
        k = gr.Slider(minimum=1, maximum=20, value=5, label="Number of results", step=1,interactive=True)
    examples = gr.Examples(
            examples,search_bar, label="Examples",
    )
    button = gr.Button("🔍 Search")
    gr.Markdown("## AI Tools")
    result_df = gr.Dataframe(
        headers=cols,
        wrap=True,
        datatype=["markdown","str","str","str","str","str"],
        column_widths = ["10%","50%","10%","10%","10%","10%"],
    )

    (button
        .click(predict, inputs = [search_bar,k], outputs=[result_df])
    )

demo.launch()