import gradio as gr
import pandas as pd

LEADERBOARD_FILE = "results.csv"

def load_leaderboard():
    return pd.read_csv(LEADERBOARD_FILE)

def display_leaderboard(score_type="Overall", open_filter="Both"):
    df = load_leaderboard()
    
    if open_filter != "Both":
        if open_filter == "Open":
            df = df[df["Open"] == True]
        elif open_filter == "Closed":
            df = df[df["Open"] == False]
    
    if score_type == "Human-Generated Split":
        df = df[["Model", "Human-Generated Safe Tasks Score", "Human-Generated Harm Tasks Score"]]
        sort_col = "Human-Generated Safe Tasks Score"
    elif score_type == "HITL Split":
        df = df[["Model", "HITL Safe Tasks Score", "HITL Harm Tasks Score"]]
        sort_col = "HITL Safe Tasks Score"
    else:
        df = df[["Model", "Overall Safe Tasks Score", "Overall Harm Tasks Score", "Refusal Rate", "Normalized Safety Score"]]
        sort_col = "Overall Safe Tasks Score"
        
    return df.sort_values(by=sort_col, ascending=False)

def display_category_scores(open_filter="Both"):
    df = load_leaderboard()
    
    if open_filter != "Both":
        if open_filter == "Open":
            df = df[df["Open"] == True]
        elif open_filter == "Closed":
            df = df[df["Open"] == False]

    df = df[["Model", "Bias", "Cybercrime", "Harassment", "Misinformation", "Illegal Activity"]]
    return df

with gr.Blocks() as demo:
    gr.Markdown("# SafeArena Leaderboard")
    
    score_filter = gr.Dropdown(label="Data Split", choices=["Full Split", "Human-Generated Split", "HITL Split"], value="Overall")
    open_state = gr.State("Both")
        
    with gr.Row():
        gr.Markdown("### Filter Models")
        open_both = gr.Button("Both", variant="primary")
        open_open = gr.Button("Open")
        open_closed = gr.Button("Closed")
        
    table_overall = gr.Dataframe(value=display_leaderboard("Overall", "Both"), interactive=False)
    table_categories = gr.Dataframe(value=display_category_scores("Both"), interactive=False)
    
    open_both.click(
        lambda score: (display_leaderboard(score, "Both"), display_category_scores("Both"), "Both"),
        inputs=[score_filter],
        outputs=[table_overall, table_categories, open_state]
    )
    open_open.click(
        lambda score: (display_leaderboard(score, "Open"), display_category_scores("Open"), "Open"),
        inputs=[score_filter],
        outputs=[table_overall, table_categories, open_state]
    )
    open_closed.click(
        lambda score: (display_leaderboard(score, "Closed"), display_category_scores("Closed"), "Closed"),
        inputs=[score_filter],
        outputs=[table_overall, table_categories, open_state]
    )
    
    score_filter.change(
        lambda score, open_val: (display_leaderboard(score, open_val), display_category_scores(open_val)),
        inputs=[score_filter, open_state],
        outputs=[table_overall, table_categories]
    )
    
    demo.launch(share=True)