from yaml import load
from persist import persist, load_widget_state
import streamlit as st

import requests
import pandas as pd

@st.cache
def get_cached_data():
    languages_df = pd.read_html("https://hf.co/languages")[0]
    languages_map = pd.Series(languages_df["Language"].values, index=languages_df["ISO code"]).to_dict()

    license_df = pd.read_html("https://huggingface.co/docs/hub/repositories-licenses")[0]
    license_map = pd.Series(
        license_df["License identifier (to use in model card)"].values, index=license_df.Fullname
    ).to_dict()

    available_metrics = [x['id'] for x in requests.get('https://huggingface.co/api/metrics').json()]

    r = requests.get('https://huggingface.co/api/models-tags-by-type')
    tags_data = r.json()
    libraries = [x['id'] for x in tags_data['library']]
    tasks = [x['id'] for x in tags_data['pipeline_tag']]
    return languages_map, license_map, available_metrics, libraries, tasks


def main():
    if "model_name" not in st.session_state:
        # Initialize session state.
        st.session_state.update({
            "model_name": "",
            "languages": [],
            "license": "",
            "library_name": "",
            "datasets": "",
            "metrics": [],
            "task": "",
            "tags": "",
            "model_description": "Some cool model...",
            "authors": "",
            "paper_url": "",
            "github_url": "",
            "bibtex_citations": "",
            "emissions": "",
        })

    languages_map, license_map, available_metrics, libraries, tasks = get_cached_data()

    st.header("Model Card Form")

    warning_placeholder = st.empty()

    st.text_input("Model Name", key=persist("model_name"))
    st.text_area("Model Description", help="The model description provides basic details about the model. This includes the architecture, version, if it was introduced in a paper, if an original implementation is available, the author, and general information about the model. Any copyright should be attributed here. General information about training procedures, parameters, and important disclaimers can also be mentioned in this section.", key=persist('model_description'))
    st.multiselect("Language(s)", list(languages_map), format_func=lambda x: languages_map[x], help="The language(s) associated with this model. If this is not a text-based model, you should specify whatever lanuage is used in the dataset. For instance, if the dataset's labels are in english, you should select English here.", key=persist("languages"))
    st.selectbox("License", [""] + list(license_map.values()), help="The license associated with this model.", key=persist("license"))
    st.selectbox("Library Name", [""] + libraries, help="The name of the library this model came from (Ex. pytorch, timm, spacy, keras, etc.). This is usually automatically detected in model repos, so it is not required.", key=persist('library_name'))
    st.text_input("Datasets (comma separated)", help="The dataset(s) used to train this model. Use dataset id from https://hf.co/datasets.", key=persist("datasets"))
    st.multiselect("Metrics", available_metrics, help="Metrics used in the training/evaluation of this model. Use metric id from https://hf.co/metrics.", key=persist("metrics"))
    st.selectbox("Task", [""] + tasks, help="What task does this model aim to solve?", key=persist('task'))
    st.text_input("Tags (comma separated)", help="Additional tags to add which will be filterable on https://hf.co/models. (Ex. image-classification, vision, resnet)", key=persist("tags"))
    st.text_input("Author(s) (comma separated)", help="The authors who developed this model. If you trained this model, the author is you.", key=persist("authors"))
    st.text_input("Related Research Paper", help="Research paper related to this model.", key=persist("paper_url"))
    st.text_input("Related GitHub Repository", help="Link to a GitHub repository used in the development of this model", key=persist("github_url"))
    st.text_area("Bibtex Citation", help="Bibtex citations for related work", key=persist("bibtex_citations"))
    st.text_input("Carbon Emitted:", help="You can estimate carbon emissions using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700)", key=persist("emissions"))

    languages=st.session_state.languages or None
    license=st.session_state.license or None
    # Handle any warnings...
    do_warn = False
    warning_msg = "Warning: The following fields are required but have not been filled in: "
    if not languages:
        warning_msg += "\n- Languages"
        do_warn = True
    if not license:
        warning_msg += "\n- License"
        do_warn = True
    if do_warn:
        warning_placeholder.error(warning_msg)
if __name__ == '__main__':
    load_widget_state()
    main()