import gradio as gr
import folium
import h3
import pandas as pd
import numpy as np
from branca.colormap import LinearColormap
from folium.plugins import MarkerCluster

def generate_sample_data():
    sf_center = (37.7749, -122.4194)
    hexagons = h3.k_ring(h3.geo_to_h3(sf_center[0], sf_center[1], 7), 20)  # Changed to resolution 7
    
    time_slots = [
        "Early Morning (5am - 8am)", "Mid-Morning (8am - 11am)", "Mid-Day (11am - 1pm)",
        "Mid-Afternoon (1pm - 4pm)", "Evening (4pm - 7pm)", "Night (7pm - 10pm)",
        "Late-Night (10pm - 1am)", "Witching Hour (1am - 4am)"
    ]
    
    data = []
    for hex_id in hexagons:
        for slot in time_slots:
            visits = np.random.poisson(100)  # Random number of visits
            data.append({"hex_id": hex_id, "time_slot": slot, "visits": visits})
    
    df = pd.DataFrame(data)
    
    # Calculate Rush Index
    avg_visits = df.groupby("time_slot")["visits"].transform("mean")
    df["rush_index"] = df["visits"] / avg_visits
    
    return df

def create_map(df, selected_time_slot):
    m = folium.Map(location=[37.7749, -122.4194], zoom_start=11, tiles="cartodbpositron")
    
    filtered_df = df[df["time_slot"] == selected_time_slot]
    
    colormap = LinearColormap(colors=['green', 'yellow', 'red'], vmin=0.5, vmax=1.5)
    
    for _, row in filtered_df.iterrows():
        hex_center = h3.h3_to_geo(row["hex_id"])
        hex_boundary = h3.h3_to_geo_boundary(row["hex_id"])
        
        color = colormap(row["rush_index"])
        
        folium.Polygon(
            locations=hex_boundary,
            color=color,
            weight=1,
            fill=True,
            fill_color=color,
            fill_opacity=0.7,
            tooltip=f"Rush Index: {row['rush_index']:.2f}"
        ).add_to(m)
    
    colormap.add_to(m)
    
    return m

def update_map(time_slot):
    map_html = create_map(df, time_slot)._repr_html_()
    return map_html

df = generate_sample_data()

time_slots = [
    "Early Morning (5am - 8am)", "Mid-Morning (8am - 11am)", "Mid-Day (11am - 1pm)",
    "Mid-Afternoon (1pm - 4pm)", "Evening (4pm - 7pm)", "Night (7pm - 10pm)",
    "Late-Night (10pm - 1am)", "Witching Hour (1am - 4am)"
]

iface = gr.Interface(
    fn=update_map,
    inputs=gr.Dropdown(choices=time_slots, label="Peak Hours"),
    outputs=gr.HTML(),
    title="RushRadar",
    description="Visualize rush hours in San Francisco using H3 hexagons",
    live=True
)

iface.launch()