import gradio as gr
import promptquality as pq
from dotenv import load_dotenv

load_dotenv()

from data_loader import (
    load_data,
    CATEGORIES,
    METHODOLOGY,
    HEADER_CONTENT,
    CARDS,
    DATASETS,
    SCORES,
)
from tabs.leaderboard import create_leaderboard_tab, filter_leaderboard
from tabs.model_comparison import create_model_comparison_tab, compare_models
from tabs.data_exploration import create_exploration_tab
from chat import filter_and_update_display


def create_app():
    df = load_data()

    MODELS = [x.strip() for x in df["Model"].unique().tolist()]

    with gr.Blocks(
        theme=gr.themes.Soft(font=[gr.themes.GoogleFont("sans-serif")])
    ) as app:
        with gr.Tabs():
            # Create tabs
            lb_output, lb_plot1, lb_plot2 = create_leaderboard_tab(
                df, CATEGORIES, METHODOLOGY, HEADER_CONTENT, CARDS
            )

            mc_info, mc_plot = create_model_comparison_tab(df, HEADER_CONTENT)

            exp_outputs = create_exploration_tab(
                df, MODELS, DATASETS, SCORES, HEADER_CONTENT
            )

        # Initial loads
        app.load(
            fn=lambda: filter_leaderboard(
                df, "All", list(CATEGORIES.keys())[0], "Performance"
            ),
            outputs=[lb_output, lb_plot1, lb_plot2],
        )

        app.load(
            fn=lambda: compare_models(
                df, [df.sort_values("Model Avg", ascending=False).iloc[0]["Model"]]
            ),
            outputs=[mc_info, mc_plot],
        )

        app.load(
            fn=lambda: filter_and_update_display(MODELS[0], DATASETS[0], 0, 1, 0),
            outputs=exp_outputs,
        )

    return app


demo = create_app()
demo.launch()