import gradio as gr
import pandas as pd
import os
import plotly.express as px
import numpy as np

datadir = 'data/emissions/complete'

model_param_df = pd.read_csv('data/model_parameters.csv', header=0)
model_performance_df = pd.read_csv('data/performance.csv', header=0)
emissions_df = pd.read_csv('data/co2_data.csv',header=0)
modalities_df = pd.read_csv('data/modalities_data.csv',header=0)
finetuned_df = emissions_df[~emissions_df['task'].str.contains('zero')]

fig0 = px.scatter(finetuned_df, x="dataset", y="query emissions (g)", color="model", log_y=True)
fig0.update_layout(xaxis={'categoryorder':'mean ascending'})
fig0.update_layout(yaxis_title='Total carbon emitted (g)')
fig0.update_layout(xaxis_title='Dataset')


fig1 = px.box(finetuned_df, x="task", y="query_energy (kWh)", color="task", log_y=True)
fig1.update_layout(xaxis={'categoryorder':'mean ascending'})
fig1.update_layout(yaxis_title='Total energy used (Wh)')
fig1.update_layout(xaxis_title='Task')

fig2 = px.scatter(modalities_df, x="num_params", y="query emissions (g)", color="modality",
             log_x=True, log_y=True, custom_data=['model','task'])

fig2.update_traces(
    hovertemplate="<br>".join([
        "Model: %{customdata[0]}",
        "Task: %{customdata[1]}",
    ])
)
fig2.update_layout(xaxis_title='Model size (number of parameters)')
fig2.update_layout(yaxis_title='Model emissions (g of CO<sub>2</sub>)')


demo = gr.Blocks()

with demo:
    gr.Markdown("# CO2 Inference Demo")
    gr.Markdown("## Explore the plots below to get more insights about the different models and tasks from our study.")
    with gr.Row():
        with gr.Column():
            gr.Plot(fig0)
    with gr.Row():
        with gr.Column():
            gr.Plot(fig1)
    with gr.Row():
        with gr.Column():
            gr.Plot(fig2)






demo.launch()