import gradio as gr
import pandas as pd

# Expanded dataframe with hidden columns
leaderboard_df = (
    pd.read_csv("results.csv", usecols=["dataset", "modality", "model", "score", "se", "metric"])
    .rename(columns={"dataset": "Dataset", "modality": "Approach", "model": "Model", "metric": "Metric"})
    .round(3)
)

def filter_leaderboard(dataset_choice, approach_choice, metric_choice):
    """Filter dataframe based on selections and drop filter columns"""
    filtered_df = leaderboard_df[
        (leaderboard_df['Dataset'] == dataset_choice) &
        (leaderboard_df['Approach'] == approach_choice) &
        (leaderboard_df['Metric'] == metric_choice)
    ]
    return filtered_df.drop(columns=['Dataset', 'Approach', 'Metric']).sort_values('score', ascending=False)

with gr.Blocks(title="TraitGym Leaderboard") as demo:
    gr.Markdown("## 🏆 TraitGym Leaderboard")
    
    with gr.Row():
        dataset_filter = gr.Radio(
            choices=["Mendelian traits", "Complex traits"],
            value="Mendelian traits",
            label="Select Dataset"
        )
        approach_filter = gr.Radio(
            choices=["Zero-shot", "Linear probing"],
            value="Zero-shot",
            label="Select Approach"
        )
        metric_filter = gr.Radio(
            choices=["AUPRC", "AUPRC_by_chrom_weighted_average"],
            value="AUPRC",
            label="Select Metric"
        )
    
    leaderboard = gr.Dataframe(
        headers=["Model", "score", "se"],
        datatype=["str", "number", "number"],
        interactive=False,
    )
    
    # Connect filters to update the table
    for filter in [dataset_filter, approach_filter, metric_filter]:
        filter.change(
            fn=filter_leaderboard,
            inputs=[dataset_filter, approach_filter, metric_filter],
            outputs=leaderboard
        )
    
    # Initial load
    demo.load(
        fn=filter_leaderboard,
        inputs=[dataset_filter, approach_filter, metric_filter],
        outputs=leaderboard
    )

demo.launch()