import gradio as gr
import pandas as pd
from huggingface_hub import list_models
import plotly.express as px

def get_plots(task_data):
    #TO DO : hover text with energy efficiency number, parameters
    task_df= pd.read_csv(task_data)
    task_df['Total GPU Energy (Wh)'] = task_df['total_gpu_energy']*1000
    task_df = task_df.sort_values(by=['Total GPU Energy (Wh)'])
    task_df['energy_star'] = pd.cut(task_df['Total GPU Energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
    fig = px.scatter(task_df, x="model", y='Total GPU Energy (Wh)', height= 500, width= 800, color = 'energy_star', color_discrete_map={"⭐": 'red', "⭐⭐": "yellow", "⭐⭐⭐": "green"})
    fig.update_traces(mode="markers+lines", hovertemplate=None)
    fig.update_layout(hovermode="y")
    return fig

def get_model_names(task_data):
    #TODO: add link to results in model card of each model
    task_df= pd.read_csv(task_data)
    model_names = task_df[['model']]
    return model_names

demo = gr.Blocks()

with demo:
    gr.Markdown(
        """# Energy Star Leaderboard

    TODO """
    )
    with gr.Tabs():
        with gr.TabItem("Text Generation 💬"):
            with gr.Row():
                with gr.Column():
                    plot = gr.Plot(get_plots('data/text_generation.csv'))
                with gr.Column():
                    table = gr.Dataframe(get_model_names('data/text_generation.csv'))

        with gr.TabItem("Image Generation 📷"):
            with gr.Row():
                with gr.Column():
                    plot = gr.Plot(get_plots('data/image_generation.csv'))
                with gr.Column():
                    table = gr.Dataframe(get_model_names('data/image_generation.csv'))

        with gr.TabItem("Text Classification 🎭"):
            with gr.Row():
                with gr.Column():
                    plot = gr.Plot(get_plots('data/text_classification.csv'))
                with gr.Column():
                    table = gr.Dataframe(get_model_names('data/text_classification.csv'))

        with gr.TabItem("Image Classification 🖼️"):
            with gr.Row():
                with gr.Column():
                    plot = gr.Plot(get_plots('data/image_classification.csv'))
                with gr.Column():
                    table = gr.Dataframe(get_model_names('data/image_classification.csv'))

        with gr.TabItem("Extractive QA ❔"):
            with gr.Row():
                with gr.Column():
                    plot = gr.Plot(get_plots('data/question_answering.csv'))
                with gr.Column():
                    table = gr.Dataframe(get_model_names('data/question_answering.csv'))

demo.launch()