import gradio as gr
import pandas as pd
import joblib
import folium

# Load the trained model and class labels
model = joblib.load('crime_model.pkl')
crime_classes = joblib.load('classes.pkl')

def predict_hotspot(lat, lon, hour, day, month):
    # Prepare input for prediction
    input_data = pd.DataFrame([[hour, day, month, lat, lon]], 
                             columns=['hour', 'day_of_week', 'month', 'latitude', 'longitude'])
    
    # Predict probabilities
    probs = model.predict_proba(input_data)[0]
    prediction = model.predict(input_data)[0]
    
    # Create a dictionary of results for the label component
    result_dict = {crime_classes[i]: float(probs[i]) for i in range(len(crime_classes))}
    
    # Generate Map
    m = folium.Map(location=[lat, lon], zoom_start=14, tiles="CartoDB positron")
    
    # Set color based on highest probability (Risk Level)
    max_prob = max(probs)
    color = "red" if max_prob > 0.4 else "orange"
    
    folium.Circle(
        location=[lat, lon],
        radius=400,
        popup=f"Predicted: {prediction}",
        color=color,
        fill=True,
        fill_opacity=0.4
    ).add_to(m)
    
    folium.Marker([lat, lon], tooltip="Query Point").add_to(m)
    
    # Convert map to HTML string
    map_html = m._repr_html_()
    
    return result_dict, map_html

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛡️ Predictive Crime Hotspot Analysis")
    gr.Markdown("Identify potential crime-prone areas using Random Forest machine learning.")
    
    with gr.Row():
        with gr.Column(scale=1):
            lat_input = gr.Number(label="Latitude", value=28.61)
            lon_input = gr.Number(label="Longitude", value=77.23)
            hour_slider = gr.Slider(0, 23, step=1, label="Hour of Day (24h)")
            day_slider = gr.Slider(0, 6, step=1, label="Day (0=Mon, 6=Sun)")
            month_slider = gr.Slider(1, 12, step=1, label="Month")
            btn = gr.Button("Analyze Risk", variant="primary")
            
        with gr.Column(scale=2):
            label_output = gr.Label(label="Crime Risk Distribution")
            map_output = gr.HTML(label="GIS Hotspot Map")

    btn.click(
        fn=predict_hotspot, 
        inputs=[lat_input, lon_input, hour_slider, day_slider, month_slider], 
        outputs=[label_output, map_output]
    )

if __name__ == "__main__":
    demo.launch()