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

def get_plots(task_df):
    grouped_df = task_df[['total_gpu_energy', 'model']].groupby('model').mean().sort_values('total_gpu_energy',ascending = False)
    grouped_df = grouped_df.reset_index()
    grouped_df['model'] = grouped_df['model'].str.split('/').str[-1]
    grouped_df['task'] = 'text_classification'
    grouped_df['total_gpu_energy (Wh)'] = grouped_df['total_gpu_energy']*1000
    grouped_df['energy_star'] = pd.cut(grouped_df['total_gpu_energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
    grouped_df = px.scatter(grouped_df, x="model", y="total_gpu_energy (Wh)", height= 500, width= 800, color = 'energy_star', color_discrete_map={"⭐": 'red', "⭐⭐": "yellow", "⭐⭐⭐": "green"})
    return grouped_df

# %% app.ipynb 3
demo = gr.Blocks()

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

    TODO """
    )
    with gr.Tabs():
        with gr.TabItem("Text Generation 💬"):
            with gr.Row():
                animal_data = gr.components.Dataframe(
                    type="pandas", datatype=["number", "markdown", "markdown", "number"]
                )
        with gr.TabItem("Image Generation 📷"):
            with gr.Row():
                science_data = gr.components.Dataframe(
                    type="pandas", datatype=["number", "markdown", "markdown", "number"]
                )
        with gr.TabItem("Text Classification 🎭"):
            with gr.Row():
                plot = gr.Plot(get_plots('data/text_classification.csv'))
        with gr.TabItem("Image Classification 🖼️"):
            with gr.Row():
                landscape_data = gr.components.Dataframe(
                    type="pandas", datatype=["number", "markdown", "markdown", "number"]
                )
        with gr.TabItem("Extractive QA ❔"):
            with gr.Row():
                wildcard_data = gr.components.Dataframe(
                    type="pandas", datatype=["number", "markdown", "markdown", "number"]
                )


demo.launch()