import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# Sample data
data = {
    "Country": ["India", "China", "USA", "Indonesia", "Brazil"],
    "Population": [1400000000, 1410000000, 331000000, 273000000, 213000000]
}
df = pd.DataFrame(data)
df.set_index("Country", inplace=True)

def generate_heatmap(selected_countries):
    filtered_df = df.loc[selected_countries]

    # Normalize for better visualization
    normalized = filtered_df.copy()
    normalized["Population"] = normalized["Population"] / 1e6  # Convert to millions

    plt.figure(figsize=(8, 4))
    sns.heatmap(normalized.T, annot=True, fmt=".1f", cmap="YlOrRd", cbar_kws={"label": "Population (in millions)"})
    plt.title("🌍 Population Heatmap")
    plt.yticks(rotation=0)
    plt.tight_layout()
    return plt

# Country options
country_choices = df.index.tolist()

# Gradio App
with gr.Blocks() as demo:
    gr.Markdown("## 🌡️ World Population Heatmap")
    selected = gr.CheckboxGroup(label="Select Countries", choices=country_choices, value=["India", "China", "USA"])
    out = gr.Plot()

    selected.change(fn=generate_heatmap, inputs=selected, outputs=out)

demo.launch()