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

# Function to load data from a given CSV file
def load_data(version):
    file_path = f'versions/{version}.csv'  # Replace with your file paths
    return pd.read_csv(file_path)

# Function for searching in the leaderboard
def search_leaderboard(df, query):
    if query == "":
        return df
    else:
        return df[df['Method'].str.contains(query)]

# Function to change the version of the leaderboard
def change_version(version):
    new_df = load_data(version)
    return new_df

# Function to create plots
from plotter import create_plots

# Initialize Gradio app
demo = gr.Blocks()

with demo:
    gr.Markdown("## 🥇 TOFU Leaderboard")

    

    with gr.Tabs():
        with gr.TabItem("Leaderboard"):
            with gr.Row():
                version_dropdown = gr.Dropdown(
                    choices=["llama", "phi", "stable-lm"],
                    label="🔄 Select Base Model",
                    value="llama",
                )

            with gr.Row():
                search_bar = gr.Textbox(
                    placeholder="Search for methods...",
                    show_label=False,
                )

            leaderboard_table = gr.components.Dataframe(
                value=load_data("llama"),
                interactive=True,
                visible=True,
            )

            version_dropdown.change(
                change_version,
                inputs=version_dropdown,
                outputs=leaderboard_table
            )

            search_bar.change(
                search_leaderboard,
                inputs=[leaderboard_table, search_bar],
                outputs=leaderboard_table
            )

        with gr.TabItem("Plots"):
            version_dropdown_plots = gr.Dropdown(
                    choices=["llama", "phi", "stable-lm"],
                    label="🔄 Select Base Model",
                    value="llama",
                )

            with gr.Row():
                methods_checkbox = gr.CheckboxGroup(
                    label="Select Methods",
                    choices=list(load_data("llama")['Method'].unique()),  # To be populated dynamically
                )

            plot_output = gr.Plot()

            # Dynamically update the choices for the methods checkbox
            def update_method_choices(version):
                df = load_data(version)
                methods = df['Method'].unique()
                methods_checkbox.update(choices=methods)
                return df

            version_dropdown_plots.change(
                update_method_choices,
                inputs=version_dropdown_plots,
                outputs=[methods_checkbox, plot_output]
            )

            methods_checkbox.change(
                create_plots,
                inputs=[methods_checkbox, leaderboard_table],
                outputs=plot_output
            )

# Launch the app
demo.launch()