import gradio as gr
import pandas as pd
import os
import fnmatch
import json


class MultiURLData:
    def __init__(self):
        self.data = self.process_data()

    
    def get_data(self):
        return self.data
    

    def process_data(self):

        dataframes = []

        def find_files(directory, pattern):
            for root, dirs, files in os.walk(directory):
                for basename in files:
                    if fnmatch.fnmatch(basename, pattern):
                        filename = os.path.join(root, basename)
                        yield filename

        for filename in find_files('results', 'results*.json'):
            model_name = filename.split('/')[2]
            with open(filename) as f:
                data = json.load(f)
                df = pd.DataFrame(data['results']).T

                # Rename 'acc' column to respective file names
                df = df.rename(columns={'acc': model_name})

                # Remove 'hendrycksTest-' from the index
                df.index = df.index.str.replace('hendrycksTest-', '')

                # Remove'harness|' from the index
                df.index = df.index.str.replace('harness\\|', '')


                dataframes.append(df[[model_name]])  # keep only the column of interest

        # Merge the dataframes on index
        data = pd.concat(dataframes, axis=1)

        # Transpose the dataframe to swap rows and columns
        data = data.transpose()
        data['Model Name'] = data.index
        cols = data.columns.tolist()
        cols = cols[-1:] + cols[:-1]
        data = data[cols]

        return data

data_provider = MultiURLData()

block = gr.Blocks()

with block:    
    gr.Markdown("""Leaderboard""")
    with gr.Tabs():
        with gr.TabItem("Leaderboard"):
            with gr.Row():
                data = gr.outputs.Dataframe(type="pandas")
            with gr.Row():
                data_run = gr.Button("Refresh")
                data_run.click(data_provider.get_data, inputs=None, outputs=data)
    # running the function on page load in addition to when the button is clicked
    block.load(data_provider.get_data, inputs=None, outputs=data)               

block.launch()