import abc

import gradio as gr

from gen_table import *
from meta_data import *

head_style = """
<style>
@media (min-width: 1536px)
{
    .gradio-container {
        min-width: var(--size-full) !important;
    }
}
</style>
"""


def math_main_tab(results):
    _, check_box = BUILD_L1_DF(results)
    table = generate_table(results)
    table['Rank'] = list(range(1, len(table) + 1))
    type_map = check_box['type_map']
    type_map['Rank'] = 'number'

    checkbox_group = gr.CheckboxGroup(choices=check_box['all'], value=check_box['required'], label='Evaluation Dimension')

    headers = ['Rank'] + check_box['essential'] + checkbox_group.value
    with gr.Row():
        model_name = gr.Textbox(value='Input the Model Name (fuzzy)', label='Model Name')
        model_size = gr.CheckboxGroup(choices=MODEL_SIZE, value=MODEL_SIZE, label='Model Size')
        model_type = gr.CheckboxGroup(choices=MODEL_TYPE, value=MODEL_TYPE, label='Model Type')

    data_component = gr.components.DataFrame(value=table[headers], datatype=[type_map[x] for x in headers])

    def filter_df(fields, model_name, model_size, model_type):
        results = load_results()['results']
        headers = ['Rank'] + check_box['essential'] + fields

        df = generate_table(results)

        df['flag'] = [model_size_flag(x, model_size) for x in df['Param (B)']]
        df = df[df['flag']]
        df.pop('flag')
        if len(df):
            df['flag'] = [model_type_flag(df.iloc[i], model_type) for i in range(len(df))]
            df = df[df['flag']]
            df.pop('flag')
        df['Rank'] = list(range(1, len(df) + 1))
        default_val = 'Input the Model Name (fuzzy)'
        if model_name != default_val:
            method_names = [x.split('</a>')[0].split('>')[-1].lower() for x in df['Method']]
            flag = [model_name.lower() in name for name in method_names]
            df['TEMP'] = flag
            df = df[df['TEMP'] == True] 
            df.pop('TEMP')

        comp = gr.components.DataFrame(value=df[headers], datatype=[type_map[x] for x in headers])
        return comp

    for cbox in [checkbox_group, model_size, model_type]:
        cbox.change(fn=filter_df, inputs=[checkbox_group, model_name, model_size, model_type], outputs=data_component)
    model_name.submit(fn=filter_df, inputs=[checkbox_group, model_name, model_size, model_type], outputs=data_component)


def dataset_tab(results, struct, dataset):
    s = struct
    s.table, s.check_box = BUILD_L2_DF(results, dataset)
    s.type_map = s.check_box['type_map']
    s.type_map['Rank'] = 'number'

    s.checkbox_group = gr.CheckboxGroup(choices=s.check_box['all'], value=s.check_box['required'], label=f'{dataset} CheckBoxes')
    s.headers = ['Rank'] + s.check_box['essential'] + s.checkbox_group.value
    s.table['Rank'] = list(range(1, len(s.table) + 1))

    with gr.Row():
        s.model_name = gr.Textbox(value='Input the Model Name (fuzzy)', label='Model Name')
        s.model_size = gr.CheckboxGroup(choices=MODEL_SIZE, value=MODEL_SIZE, label='Model Size')
        s.model_type = gr.CheckboxGroup(choices=MODEL_TYPE, value=MODEL_TYPE, label='Model Type')

    s.data_component = gr.components.DataFrame(value=s.table[s.headers], datatype=[s.type_map[x] for x in s.headers])
    s.dataset = gr.Textbox(value=dataset, label=dataset, visible=False)

    def filter_df_l2(dataset_name, fields, model_name, model_size, model_type):
        results = load_results()['results']
        s = structs[DATASETS.index(dataset_name)]
        headers = ['Rank'] + s.check_box['essential'] + fields
        df = cp.deepcopy(s.table)
        df['flag'] = [model_size_flag(x, model_size) for x in df['Param (B)']]
        df = df[df['flag']]
        df.pop('flag')
        if len(df):
            df['flag'] = [model_type_flag(df.iloc[i], model_type) for i in range(len(df))]
            df = df[df['flag']]
            df.pop('flag')
        df['Rank'] = list(range(1, len(df) + 1))
        default_val = 'Input the Model Name (fuzzy)'
        if model_name != default_val:
            method_names = [x.split('</a>')[0].split('>')[-1].lower() for x in df['Method']]
            flag = [model_name.lower() in name for name in method_names]
            df['TEMP'] = flag
            df = df[df['TEMP'] == True] 
            df.pop('TEMP')

        comp = gr.components.DataFrame(value=df[headers], datatype=[s.type_map[x] for x in headers])
        return comp

    for cbox in [s.checkbox_group, s.model_size, s.model_type]:
        cbox.change(
            fn=filter_df_l2,
            inputs=[s.dataset, s.checkbox_group, s.model_name, s.model_size, s.model_type],
            outputs=s.data_component)
    s.model_name.submit(
        fn=filter_df_l2, 
        inputs=[s.dataset, s.checkbox_group, s.model_name, s.model_size, s.model_type],
        outputs=s.data_component)

    
with gr.Blocks(title="Math Leaderboard", head=head_style) as demo:
    results = load_results()['results']
    N_MODEL = len(results)
    DATASETS = []
    for m in results:
        DATASETS.extend(results[m].keys())
    DATASETS = [d for d in set(DATASETS) if d != 'META']
    

    N_DATA = len(DATASETS)
    structs = [abc.abstractproperty() for _ in range(N_DATA)]

    gr.Markdown(LEADERBORAD_INTRODUCTION)

    with gr.Tabs(elem_classes='tab-buttons') as tabs:
        with gr.TabItem('🏅 LMM Math Leaderboard', elem_id='main', id=0):
            math_main_tab(results)
            
        for i, dataset in enumerate(DATASETS):
            tab_name_map = {
                'MathVista': 'MathVista (Test Mini)',
                'MathVerse': 'MathVerse (Vision Only)',
            }

            with gr.TabItem(
                f'📊 {dataset if dataset not in tab_name_map else tab_name_map[dataset]}', elem_id=dataset, id=i + 2):
                dataset_tab(results, structs[i], dataset)

    with gr.Row():
        with gr.Accordion('Citation', open=False):
            citation_button = gr.Textbox(
                value=CITATION_BUTTON_TEXT,
                label=CITATION_BUTTON_LABEL,
                elem_id='citation-button')

if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0')