import gradio as gr
import plotly.graph_objects as go
from datasets import load_dataset

# Load dataset once
dataset = load_dataset("gradio/NYC-Airbnb-Open-Data", split="train")
df = dataset.to_pandas()

def filter_map(min_price, max_price, boroughs):
    # Handle empty boroughs list
    if not boroughs:
        boroughs = ["Queens", "Brooklyn", "Manhattan", "Bronx", "Staten Island"]
    
    # Filter dataframe
    filtered_df = df[(df['neighbourhood_group'].isin(boroughs)) &
                    (df['price'] > min_price) & (df['price'] < max_price)]
    
    # Handle empty results
    if filtered_df.empty:
        # Return empty map
        fig = go.Figure()
        fig.update_layout(
            title="No properties found with current filters",
            mapbox_style="open-street-map",
            mapbox=dict(
                center=go.layout.mapbox.Center(lat=40.67, lon=-73.90),
                zoom=9
            ),
        )
        return fig
    
    # Prepare data for map
    names = filtered_df["name"].tolist()
    prices = filtered_df["price"].tolist()
    text_list = [(names[i], prices[i]) for i in range(len(names))]
    
    # Create map
    fig = go.Figure(go.Scattermapbox(
        customdata=text_list,
        lat=filtered_df['latitude'].tolist(),
        lon=filtered_df['longitude'].tolist(),
        mode='markers',
        marker=go.scattermapbox.Marker(
            size=6,
            color='red',
            opacity=0.7
        ),
        hoverinfo="text",
        hovertemplate='<b>Name</b>: %{customdata[0]}<br><b>Price</b>: $%{customdata[1]}<extra></extra>'
    ))

    fig.update_layout(
        title=f"Found {len(filtered_df)} properties",
        mapbox_style="open-street-map",
        hovermode='closest',
        mapbox=dict(
            bearing=0,
            center=go.layout.mapbox.Center(
                lat=40.67,
                lon=-73.90
            ),
            pitch=0,
            zoom=9
        ),
        height=600
    )

    return fig

with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            min_price = gr.Number(value=250, label="Minimum Price")
            max_price = gr.Number(value=1000, label="Maximum Price")
        boroughs = gr.CheckboxGroup(choices=["Queens", "Brooklyn", "Manhattan", "Bronx", "Staten Island"], value=["Queens", "Brooklyn"], label="Select Boroughs:")
        btn = gr.Button(value="Update Filter")
        map = gr.Plot()
    demo.load(filter_map, [min_price, max_price, boroughs], map)
    btn.click(filter_map, [min_price, max_price, boroughs], map)

demo.launch()