import gradio as gr
import pandas as pd

from themes import Seafoam
from load_data import load_main_table
from constants import BANNER, CITATION_TEXT, css, js_code, all_task_types, js_light

TYPES = ["number", "markdown", "number"]
MAIN_TABLE_COLS = ['Model', 'Language', 'Average Toxicity', 'Expected Maximum Toxicity', 'Empirical Probability']



df_main = load_main_table()
available_models = df_main['Model'].unique()
MODEL_SIZE = list(df_main['Model Size'].unique())
MODEL_TYPE = list(df_main['Model Type'].unique())
LANGAUGES = list(df_main['Language'].unique())
MODEL_FAMILY = list(df_main['Model Family'].unique())


with open("_about_us.md", "r") as f:
    ABOUT_MD = f.read()

with open("_header.md", "r") as f:
    HEADER_MD = f.read()

with open("_metrics.md", "r") as f:
    METRIC_MD = f.read()

    
with gr.Blocks(theme=gr.themes.Soft(), css=css, js=js_light) as demo:
    gr.HTML(BANNER, elem_id="banner")
    # gr.HTML("<img src='file/image.png' alt='image One'>")
    gr.Markdown(HEADER_MD, elem_classes="markdown-text")
    # gr.Image("./data/ptp.png")
    gr.HTML("<img src='file/data/ptp.png' alt='image One'>")
    gr.Markdown(f"**Version**: PTP-Small | **# Examples**: 85K | **# Models**: {len(available_models)}", elem_classes="markdown-text")
    gr.Markdown(METRIC_MD, elem_classes="markdown-text")


    with gr.Tabs(elem_classes="tab-buttons") as tabs:

        with gr.TabItem("🏅 Multilingual Leaderboard", elem_id="od-benchmark-tab-table", id=0, elem_classes="subtab"):
            print(df_main.head())
            mling_df = df_main.loc[df_main['Multilingual']==True, MAIN_TABLE_COLS].copy()
            del mling_df['Language']
            
            mling_df = mling_df.groupby("Model").agg('mean').reset_index().round(3)
            mling_df = mling_df.sort_values(by="Average Toxicity")

            print(mling_df.head())

            ablation_table = gr.components.Dataframe(
                        value=mling_df,
                        datatype=TYPES,
                        height=1000,
                        elem_id="mling-table",
                        interactive=False,
                        visible=True,
                        min_width=60,
            )


        with gr.TabItem("📊 Ablation Results", elem_id="od-benchmark-tab-table", id=1, elem_classes="subtab"):
            with gr.Row():
                language = gr.CheckboxGroup(
                    choices=LANGAUGES,
                    value=LANGAUGES,
                    label='Language',
                    interactive=True
                )
            with gr.Row():
                model_family = gr.CheckboxGroup(
                    choices=MODEL_FAMILY,
                    value=MODEL_FAMILY,
                    label='Model Family',
                    interactive=True
                )
                
            with gr.Row():
                model_size = gr.CheckboxGroup(
                    choices=MODEL_SIZE,
                    value=MODEL_SIZE,
                    label='Model Size',
                    interactive=True,
                )
                model_type = gr.CheckboxGroup(
                    choices=MODEL_TYPE,
                    value=MODEL_TYPE,
                    label='Model Type',
                    interactive=True
                )

            ablation_table = gr.components.Dataframe(
                        value=df_main[MAIN_TABLE_COLS],
                        datatype=TYPES,
                        height=500,
                        elem_id="full-table",
                        interactive=False,
                        visible=True,
                        min_width=60,
            )    


            def filter_df(model_size, model_type, language, model_family):
                df = df_main.copy()
                
                print(df.isnull().sum())

                df = df[df['Model Type'].isin(model_type)]
                df = df[df['Model Size'].isin(model_size)]
                df = df[df['Language'].isin(language)]
                df = df[df['Model Family'].isin(model_family)]

                df = df.sort_values(by="Average Toxicity")

                assert (df.isnull().sum().sum())==0


                comp = gr.components.DataFrame(
                    value=df[MAIN_TABLE_COLS],
                    datatype=TYPES,
                    interactive=False,
                    visible=True)
                return comp
            
            for cbox in [model_size, model_type, language, model_family]:
                cbox.change(fn=filter_df, inputs=[model_size, model_type, language, model_family], outputs=ablation_table)


        with gr.TabItem("📮 About Us", elem_id="od-benchmark-tab-table", id=2):
                gr.Markdown(ABOUT_MD, elem_classes="markdown-text")
        
        with gr.Row():
            with gr.Accordion("📙 Citation", open=False, elem_classes="accordion-label"):
                gr.Textbox(
                    value=CITATION_TEXT, 
                    lines=7,
                    label="Copy the BibTeX snippet to cite this source",
                    elem_id="citation-button",
                    show_copy_button=True)

if __name__ == '__main__':
    demo.launch(share=True, allowed_paths=["."])