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>
"""

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):
            _, 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',
                interactive=True,
            )

            headers = ['Rank'] + check_box['essential'] + checkbox_group.value
            with gr.Row():
                model_name = gr.Textbox(
                    value='Input the Model Name (fuzzy)', 
                    label='Model Name', 
                    interactive=True,
                    visible=True)
                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
                )

            data_component = gr.components.DataFrame(
                value=table[headers],
                type='pandas',
                datatype=[type_map[x] for x in headers],
                interactive=False,
                visible=True)

            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:
                    print(model_name)
                    model_name = model_name.lower()
                    method_names = [x.split('</a>')[0].split('>')[-1].lower() for x in df['Method']]
                    flag = [model_name in name for name in method_names]
                    df['TEMP_FLAG'] = flag
                    df = df[df['TEMP_FLAG'] == True] 
                    df.pop('TEMP_FLAG')

                comp = gr.components.DataFrame(
                    value=df[headers],
                    type='pandas',
                    datatype=[type_map[x] for x in headers],
                    interactive=False,
                    visible=True)
                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)

        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):

                s = structs[i]
                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',
                    interactive=True,
                )
                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', 
                        interactive=True,
                        visible=True)
                    s.model_size = gr.CheckboxGroup(
                        choices=MODEL_SIZE,
                        value=MODEL_SIZE,
                        label='Model Size',
                        interactive=True
                    )
                    s.model_type = gr.CheckboxGroup(
                        choices=MODEL_TYPE,
                        value=MODEL_TYPE,
                        label='Model Type',
                        interactive=True
                    )
                s.data_component = gr.components.DataFrame(
                    value=s.table[s.headers],
                    type='pandas',
                    datatype=[s.type_map[x] for x in s.headers],
                    interactive=False,
                    visible=True)
                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:
                        print(model_name)
                        model_name = model_name.lower()
                        method_names = [x.split('</a>')[0].split('>')[-1].lower() for x in df['Method']]
                        flag = [model_name in name for name in method_names]
                        df['TEMP_FLAG'] = flag
                        df = df[df['TEMP_FLAG'] == True] 
                        df.pop('TEMP_FLAG')

                    comp = gr.components.DataFrame(
                        value=df[headers],
                        type='pandas',
                        datatype=[s.type_map[x] for x in headers],
                        interactive=False,
                        visible=True)
                    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.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')