import pandas as pd
import re
from huggingface_hub import HfApi
api = HfApi()

def get_model_size(model_name, precision: str = "BF16", revision: str = "main"):
    model_info = api.model_info(repo_id=model_name, revision=revision)
    # model_size = get_model_size(model_info=model_info, precision=precision)
    size_pattern = size_pattern = re.compile(r"(\d\.)?\d+(b|m)")
    try:
        model_size = round(model_info.safetensors["total"] / 1e9, 3)
    except (AttributeError, TypeError):
        try:
            size_match = re.search(size_pattern, model_info.modelId.lower())
            model_size = size_match.group(0)
            model_size = round(float(model_size[:-1]) if model_size[-1] == "b" else float(model_size[:-1]) / 1e3, 3)
        except AttributeError:
            return 0  # Unknown model sizes are indicated as 0, see NUMERIC_INTERVALS in app.py

    size_factor = 8 if (precision == "GPTQ" or "gptq" in model_info.modelId.lower()) else 1
    model_size = size_factor * model_size
    return model_size

def load_data(data_path):
    df = pd.read_csv(data_path, skiprows=1, header=0).dropna()

    columns = ['Model', 'type', 'open?', 'shot', 'en', 'zh', 'id', 'th', 'vi', 'avg', 'avg_sea']
    columns_sorted = ['rank','type', 'Model', 'open?', 'avg_sea', 'en', 'zh', 'id', 'th', 'vi', 'avg']

    # Splitting into three separate DataFrames based on the groups M3Exam and MMLU and average
    df_m3exam = df.iloc[:, :11]  # M3Exam columns
    df_mmlu = df.iloc[:, [0, 1, 2, 3, 11, 12, 13, 14, 15, 16, 17]]  # MMLU columns
    df_avg = df.iloc[:, [0, 1, 2, 3, 18, 19, 20, 21, 22, 23, 24]]  # Average columns
    df_mmlu.columns = columns
    df_avg.columns = columns

    # # multiply the values in the ['en', 'zh', 'id', 'th', 'vi', 'avg', 'avg_sea'] by 100 and display as 1 decimal
    for df_tmp in [df_m3exam, df_mmlu, df_avg]:
        df_tmp[['en', 'zh', 'id', 'th', 'vi', 'avg', 'avg_sea']] *= 100
        df_tmp[['en', 'zh', 'id', 'th', 'vi', 'avg', 'avg_sea']] = df_tmp[['en', 'zh', 'id', 'th', 'vi', 'avg', 'avg_sea']].round(2)

    # rank the DataFrames by the 'avg_sea' column
    df_m3exam['rank'] = df_m3exam['avg_sea'].rank(ascending=False).astype(int)
    df_mmlu['rank'] = df_mmlu['avg_sea'].rank(ascending=False).astype(int)
    df_avg['rank'] = df_avg['avg_sea'].rank(ascending=False).astype(int)

    # reorder the columns
    df_m3exam = df_m3exam[columns_sorted]
    df_mmlu = df_mmlu[columns_sorted]
    df_avg = df_avg[columns_sorted]

    # sort the DataFrames by the 'avg_sea' column in descending order
    df_m3exam = df_m3exam.sort_values(by='avg_sea', ascending=False)
    df_mmlu = df_mmlu.sort_values(by='avg_sea', ascending=False)
    df_avg = df_avg.sort_values(by='avg_sea', ascending=False)

    # change the column name from 'avg_sea' to 'avg_sea ⬇️'
    df_m3exam = df_m3exam.rename(columns={'avg_sea': 'avg_sea ⬇️'})
    df_mmlu = df_mmlu.rename(columns={'avg_sea': 'avg_sea ⬇️'})
    df_avg = df_avg.rename(columns={'avg_sea': 'avg_sea ⬇️'})

    # map the values in the 'type' column to the following values: {'base': 'Base', 'chat': 'Chat'}
    df_m3exam['type'] = df_m3exam['type'].map({'base': '🟢 base', 'chat': '🔶 chat'})
    df_mmlu['type'] = df_mmlu['type'].map({'base': '🟢 base', 'chat': '🔶 chat'})
    df_avg['type'] = df_avg['type'].map({'base': '🟢 base', 'chat': '🔶 chat'})

    return df_m3exam, df_mmlu, df_avg


if __name__ == "__main__":
    model_name = "SeaLLMs/SeaLLM-7B-v2"
    model_size = get_model_size(model_name)
    print(model_size)