import gradio as gr
import pandas as pd
import altair as alt

# 加载数据
df = pd.read_csv("data.csv")

def plot_income_share(countries, year_range):
    """根据选中的国家和年份范围，生成图表"""
    # 过滤数据
    filtered = df[df["Entity"].isin(countries)]
    filtered = filtered[(filtered["Year"] >= year_range[0]) & 
                        (filtered["Year"] <= year_range[1])]
    
    if filtered.empty:
        return "<p style='text-align: center; color: gray;'>No data available for selected criteria.</p>"
    
    # 创建 Altair 图表
    chart = alt.Chart(filtered).mark_line().encode(
        x=alt.X("Year", axis=alt.Axis(format="d", title="Year")),
        y=alt.Y("Percent", axis=alt.Axis(title="Percent", format="~s")),
        color=alt.Color("Entity", legend=alt.Legend(title="Country")),
        strokeDash="Entity",
    ).properties(
        title="Top 5% Income Share Over Time",
        height=500
    ).interactive()
    
    return chart.to_html()

def show_raw_data(countries, year_range):
    """显示原始数据表格"""
    filtered = df[df["Entity"].isin(countries)]
    filtered = filtered[(filtered["Year"] >= year_range[0]) & 
                        (filtered["Year"] <= year_range[1])]
    return filtered

# 获取所有国家列表和年份范围
all_countries = sorted(df["Entity"].unique())
min_year, max_year = int(df["Year"].min()), int(df["Year"].max())

# 创建 Gradio 界面
with gr.Blocks(title="Top 5% Income Share", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📊 Top 5% Income Share")
    gr.Markdown("Share of income received by the richest 5% of the population. Data source: Our World in Data")
    
    with gr.Row():
        with gr.Column(scale=1):
            countries_input = gr.Dropdown(
                choices=all_countries,
                value=["Australia", "China", "Germany", "Japan", "United States"],
                multiselect=True,
                label="🌍 Select Countries"
            )
            year_input = gr.Slider(
                minimum=min_year,
                maximum=max_year,
                value=[min_year, max_year],
                step=1,
                label="📅 Year Range"
            )
            submit_btn = gr.Button("🔄 Update Chart", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            chart_output = gr.HTML(label="Income Share Chart")
    
    with gr.Accordion("📋 Show Raw Data", open=False):
        raw_data_output = gr.Dataframe(label="Filtered Data", height=300)
    
    # 绑定事件
    submit_btn.click(
        fn=plot_income_share,
        inputs=[countries_input, year_input],
        outputs=chart_output
    )
    
    submit_btn.click(
        fn=show_raw_data,
        inputs=[countries_input, year_input],
        outputs=raw_data_output
    )
    
    # 页面加载时自动显示默认数据
    demo.load(
        fn=plot_income_share,
        inputs=[countries_input, year_input],
        outputs=chart_output
    )
    demo.load(
        fn=show_raw_data,
        inputs=[countries_input, year_input],
        outputs=raw_data_output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)