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

# Load the data
df = pd.read_csv("data.csv")

# Get unique entities and years
entities = sorted(df["Entity"].unique().tolist())
earliest_year = int(df["Year"].min())
latest_year = int(df["Year"].max())

def plot_data(selected_entities, start_year, end_year, show_raw_data):
    if not selected_entities:
        return None, pd.DataFrame()

    filtered_df = df[df["Entity"].isin(selected_entities)].copy()
    filtered_df = filtered_df[(filtered_df["Year"] >= start_year) & (filtered_df["Year"] <= end_year)]

    # Create the plot
    fig, ax = plt.subplots(figsize=(10, 6))
    for entity in selected_entities:
        entity_data = filtered_df[filtered_df["Entity"] == entity]
        if not entity_data.empty:
            ax.plot(entity_data["Year"], entity_data["Percent"], label=entity, marker="o", markersize=4)

    ax.set_xlabel("Year")
    ax.set_ylabel("Income Share (%)")
    ax.set_title("Top 5% Income Share Trend")
    ax.legend(title="Country", bbox_to_anchor=(1.05, 1), loc="upper left")
    ax.grid(True, linestyle="--", alpha=0.6)
    plt.tight_layout()

    if show_raw_data:
        return fig, filtered_df
    else:
        return fig, pd.DataFrame()

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 📊 Top 5% Income Share Visualization")
    
    with gr.Row():
        with gr.Column():
            entity_selector = gr.CheckboxGroup(
                choices=entities,
                label="Select Countries",
                value=["Australia", "China", "Germany", "Japan", "United States"]
            )
            with gr.Row():
                start_year_input = gr.Number(value=earliest_year, label="Start Year", precision=0)
                end_year_input = gr.Number(value=latest_year, label="End Year", precision=0)
            
            show_raw_data_checkbox = gr.Checkbox(label="Show Data Table", value=False)
            submit_button = gr.Button("Update Chart", variant="primary")

        with gr.Column():
            output_plot = gr.Plot()
            output_dataframe = gr.DataFrame()

    submit_button.click(
        fn=plot_data,
        inputs=[entity_selector, start_year_input, end_year_input, show_raw_data_checkbox],
        outputs=[output_plot, output_dataframe]
    )

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