import gradio as gr
import pandas as pd
import numpy as np

# Read the CSV file
df = pd.read_csv('passk.csv')

# Extract unique languages and models from the Dataset column
languages = sorted(set([d.split('-')[1] for d in df['Dataset']]))
models = sorted(set(['-'.join(d.split('-')[2:-2]) for d in df['Dataset']]))

# Create a dictionary to map models to friendly names
model_to_friendly = {
    "starcoder2_15b": "StarCoder2-15B",
    "deepseekcoder_v2lite": "DeepSeekCoder2-Lite"
}

# Function to get friendly name or original name if not in the dictionary
def get_friendly_name(model):
    return model_to_friendly.get(model, model)

# Create a pivot table
pivot = df.pivot(index='Dataset', columns='Dataset', values='Estimate')
pivot.index = ['-'.join(i.split('-')[2:-2]) for i in pivot.index]
pivot.columns = [i.split('-')[1] for i in pivot.columns]

# Function to update the table based on selected languages
def update_table(selected_languages):
    if not selected_languages:
        return np.full((len(models), len(languages)), "-").tolist()
    
    display_data = pivot[selected_languages].replace(np.nan, "-")
    display_data = display_data.applymap(lambda x: f"{x:.3f}" if isinstance(x, (int, float)) else x)
    display_data.index = [get_friendly_name(model) for model in display_data.index]
    return display_data.values.tolist()

# Create the Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# Model Leaderboard")
    
    with gr.Row():
        language_checkboxes = gr.CheckboxGroup(choices=languages, label="Select Languages", value=languages)
    
    table = gr.Dataframe(
        headers=[lang.capitalize() for lang in languages],
        row_headers=[get_friendly_name(model) for model in models],
        col_count=(lambda: len(languages)),
        interactive=False
    )
    
    language_checkboxes.change(update_table, inputs=[language_checkboxes], outputs=[table])

# Launch the app
if __name__ == "__main__":
    app.launch()