import pandas as pd
import yaml
import numpy as np

def add_model(metadata_archive):
    """
    Esto actualiza el archivo del cual app.py coge la información para crear la leaderboard. Entonces, cuando
    alguien quiera añadir un nuevo modelo, tiene que ejecutar este archivo. 

    1. Leemos el CSV, sacamos información y añadimos simplemente una nueva row. 
    
    """
    # Initialize an empty DataFrame
    df = pd.DataFrame(columns=['dataset_name', 'Accuracy', 'Spearman', "Category"])

    metadata_archive = 'mteb_metadata.yaml'

    with open(metadata_archive, 'r') as file:
        for index, data in enumerate(yaml.safe_load_all(file)):
            if index == 0:
                model_index_list = data.get('model-index', [])
                model_name = model_index_list[0].get('name')
                results_list = model_index_list[0].get('results', [])
            
                if results_list:  
                    for i in range(len(results_list)):
                        task = results_list[i].get('task', {})
                        task_name = task.get("type")
                        dataset_name = results_list[i]['dataset']['name']
                        
                        # Initialize the row with NaN values
                        row = {'dataset_name': dataset_name, 'Accuracy': None, 'Spearman': None}
                        
                        if task_name == "Classification":
                            accuracy = next((metric.get('value') for metric in results_list[i].get('metrics', []) if metric.get('type') == 'accuracy'), None)
                            row['Accuracy'] = accuracy
                            row['Category'] = "Classification"
                        elif task_name == "STS":
                            spearman = next((metric.get('value') for metric in results_list[i].get('metrics', []) if metric.get('type') == 'cos_sim_spearman'), None)
                            row['Spearman'] = spearman
                            row["Category"] = "STS"
                        
                        # Append the row to the DataFrame using pd.concat
                        new_df = pd.DataFrame([row])
                        df = pd.concat([df, new_df], ignore_index=True)

    df['Accuracy'] = pd.to_numeric(df['Accuracy'], errors='coerce')
    classification_average = round(df.loc[df['Category'] == 'Classification', 'Accuracy'].mean(),2)

    df['Spearman'] = pd.to_numeric(df['Spearman'], errors='coerce')
    sts_spearman_average = round(df.loc[df['Category'] == 'STS', 'Spearman'].mean(),2)

    ## CLASSIFICATION
    classification_dataframe = pd.read_csv('../data/classification.csv')
    classification_df = df[df['Category']== 'Classification']
    new_row_data = {'Model name': model_name, 'Average': classification_average}
    for index, row in classification_df.iterrows():
        column_name = row['dataset_name']
        accuracy_value = row['Accuracy']
        new_row_data[column_name] = round(accuracy_value,2)
    
    new_row_df = pd.DataFrame(new_row_data,index=[0])
    classification_dataframe = pd.concat([classification_dataframe,new_row_df],ignore_index=True)
    classification_dataframe.to_csv("../data/classification.csv",index=False)
    
    ## STS
    sts_dataframe = pd.read_csv('../data/sts.csv')
    sts_df = df[df['Category']=='STS']
    new_row_data = {'Model name': model_name, 'Average': sts_spearman_average}

    for index, row in sts_df.iterrows():
        column_name = row['dataset_name']
        spearman_value = row['Spearman']
        new_row_data[column_name] = round(spearman_value,2)
    
    new_row_df = pd.DataFrame(new_row_data,index = [0])
    sts_dataframe = pd.concat([sts_dataframe,new_row_df],ignore_index=True)
    sts_dataframe.to_csv('../data/sts.csv',index=False)

    ## GENERAL
    general_dataframe = pd.read_csv("../data/general.csv")

    average = round(np.mean([classification_average,sts_spearman_average]),2)
    ## TODO: solucionar la meta-data como Model Size o Embedding Dimensions. 
    new_instance = {'Model name':model_name, 'Model Size (GB)': None, 'Embedding Dimensions': None, 'Average':average, 'Classification Average': classification_average, 'Clustering Average': None, 'STS Average': sts_spearman_average, 'Retrieval Average': None}
    new_row_df = pd.DataFrame(new_instance, index=[0])
    general_dataframe = pd.concat([general_dataframe, new_row_df], ignore_index=True)
    general_dataframe.to_csv("../data/general.csv",index=False)

add_model('mteb_metadata.yaml')