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"):
    if len(model_name.split("/")) == 1:
        return None

    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, 1)
    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, 1)
        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 make_clickable_model(model_name, link=None):
    if len(model_name.split("/")) == 2:
        link = "https://huggingface.co/" + model_name
        return (
            f'<a target="_blank" style="text-decoration: underline" href="{link}">{model_name.split("/")[-1]}</a>'
        )
    return model_name


def load_data(data_path):
    columns = ['Rank', 'Model', 'From', 'Open?', 'Params(B)', 'Cost', 'Score']
    columns_sorted = ['Rank', 'Model', 'From', 'Open?', 'Params(B)', 'Cost', 'Score']
    
    df = pd.read_csv(data_path, usecols=columns).dropna()
    df['Score'] = df['Score'].round(0)

    # rank according to the Score column
    df = df.sort_values(by='Score', ascending=False)
    # reorder the columns
    df = df[columns_sorted]

    # make the 'Model' column clickable
    df['Model'] = df['Model'].apply(make_clickable_model)

    return df


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