from map_search import query_data
import gradio as gr
from gradio.themes.base import Base
import pandas as pd
import plotly.graph_objects as go

def query_data_with_map(question):
    if question:
        # Assuming query_data returns a tuple: (list_of_arrays, text_from_gemini)
        list_of_arrays, text_from_gemini = query_data(question)
        
        # Convert list_of_arrays to a DataFrame
        df = pd.DataFrame(list_of_arrays, columns=['title', 'chalet_title', 'final_price', 'unit_custom_title', 'lat', 'lng'])
        
        text_list = [(row['title'], row['unit_custom_title'], row['chalet_title'], row['final_price']) for index, row in df.iterrows()]
        fig = go.Figure(go.Scattermapbox(
            customdata=text_list,
            lat=df['lat'].tolist(),
            lon=df['lng'].tolist(),
            mode='markers',
            marker=go.scattermapbox.Marker(
                size=6
            ),
            hoverinfo="text",
            hovertemplate=(
                '<b>Title</b>: %{customdata[0]}<br>'
                '<b>Unit Custom Title</b>: %{customdata[1]}<br>'
                '<b>Chalet Title</b>: %{customdata[2]}<br>'
                '<b>Final Price</b>: SAR %{customdata[3]}'
            )
        ))
    else:
        # Create an empty map
        fig = go.Figure(go.Scattermapbox(
            lat=[],
            lon=[],
            mode='markers',
            marker=go.scattermapbox.Marker(
                size=20
            )
        ))
        text_from_gemini = ""

    fig.update_layout(
        mapbox_style="open-street-map",
        hovermode='closest',
        mapbox=dict(
            bearing=0,
            center=go.layout.mapbox.Center(
                lat=24.7136,  # Latitude for Riyadh
                lon=46.6753   # Longitude for Riyadh
            ),
            pitch=0,
            zoom=10
        )
    )
    return fig, text_from_gemini

with gr.Blocks(theme=Base(), title="Riyadh Entertainment Map powered by Smart Search System using Vector Search + RAG") as demo:
    gr.Markdown(
        """
        # Smart Search System using Atlas Vector Search + RAG Architecture
        """)
    textbox = gr.Textbox(label="Enter your query here", lines=1)
    with gr.Row():
        button = gr.Button("Search", variant="primary")
    with gr.Column():
        output1 = gr.Plot(label="Map Output")
        output2 = gr.Textbox(lines=1, max_lines=10, label="Output generated by chaining Atlas Vector Search to Langchain's `load_qa_chain` + Gemini flash 1.5 LLM:")

    # Load the empty map when the app starts
    demo.load(query_data_with_map, inputs=[textbox], outputs=[output1, output2])
    # Call query_data_with_map function upon clicking the Submit button
    button.click(query_data_with_map, textbox, outputs=[output1, output2])

demo.launch(share=True)