import gradio as gr
import pandas as pd
import plotly.express as px
import matplotlib.pyplot as plt

# Sample data with coordinates (approximate)
data = {
    "Country": ["India", "China", "USA", "Indonesia", "Brazil"],
    "Population": [1400000000, 1410000000, 331000000, 273000000, 213000000],
    "Latitude": [20.5937, 35.8617, 37.0902, -0.7893, -14.2350],
    "Longitude": [78.9629, 104.1954, -95.7129, 113.9213, -51.9253]
}
df = pd.DataFrame(data)

# Geo heatmap function
def generate_geo_heatmap(selected_countries):
    filtered_df = df[df["Country"].isin(selected_countries)]

    fig = px.density_mapbox(
        filtered_df,
        lat="Latitude",
        lon="Longitude",
        z="Population",
        hover_name="Country",
        radius=30,
        center=dict(lat=20, lon=0),
        zoom=1,
        mapbox_style="carto-positron"
    )
    fig.update_layout(title="🌍 World Population Heatmap (Globe Map)", height=500)
    return fig

# Pie chart function (same as before)
def generate_pie_chart(selected_countries):
    filtered_df = df[df["Country"].isin(selected_countries)]
    plt.figure(figsize=(5, 5))
    plt.pie(filtered_df["Population"], labels=filtered_df["Country"], autopct="%1.1f%%", startangle=90)
    plt.title("🧩 Population Distribution")
    plt.tight_layout()
    return plt

# Gradio App
with gr.Blocks() as demo:
    gr.Markdown("## 🌍 World Population Heatmap and Pie Chart")
    selected = gr.CheckboxGroup(label="Select Countries", choices=df["Country"].tolist(), value=["India", "China", "USA"])

    with gr.Row():
        heatmap_output = gr.Plot(label="World Heatmap (Map View)")
        piechart_output = gr.Plot(label="Pie Chart")

    selected.change(fn=generate_geo_heatmap, inputs=selected, outputs=heatmap_output)
    selected.change(fn=generate_pie_chart, inputs=selected, outputs=piechart_output)

demo.launch()