import streamlit as st

from huggingface_hub import HfApi
import pandas
from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES

hf_api = HfApi()
all_stats = {}
total_downloads = 0

for model_name in list(CONFIG_MAPPING_NAMES.keys())[:2]:
    model_stats = {"num_downloads": 0, "%_of_all_downloads": 0, "num_models": 0, "download_per_model": 0}
    models = hf_api.list_models(filter=model_name)

    model_stats["num_models"] = len(models)
    model_stats["num_downloads"] = sum([m.downloads for m in models if hasattr(m, "downloads")])
    if len(models) > 0:
        model_stats["download_per_model"] = round(model_stats["num_downloads"] / len(models), 2)
    total_downloads += model_stats["num_downloads"]

    # save in overall dict
    all_stats[model_name] = model_stats


for model_name in list(CONFIG_MAPPING_NAMES.keys()):
    all_stats[model_name]["%_of_all_downloads"] = round(all_stats[model_name]["num_downloads"] / total_downloads, 5) * 100  # noqa: E501
    downloads = all_stats[model_name]["num_downloads"]
    all_stats[model_name]["num_downloads"] = f"{downloads:,}"


sorted_results = dict(reversed(sorted(all_stats.items(), key=lambda d: d[1]["%_of_all_downloads"])))
dataframe = pandas.DataFrame.from_dict(sorted_results, orient="index")
result = dataframe.to_string()

with open("result.txt", "w") as f:
    f.write(result)

st.table(dataframe)