import gradio as gr
import pandas as pd
from hub_utils import check_for_discussion, report_results
from model_utils import calculate_memory, get_model
from huggingface_hub.utils import HfHubHTTPError
from hub_model_stats_utils import get_model_type_downloads

# We need to store them as globals because gradio doesn't have a way for us to pass them into the button
MODEL = None
TASK_INP = None

def get_mem_results(model_name: str, library: str, options: list, access_token: str):
    global MODEL
    MODEL = get_model(model_name, library, access_token)
    try:
        has_discussion = check_for_discussion(model_name)
    except HfHubHTTPError:
        has_discussion = True
    title = f"## Memory usage for '{model_name}'"
    data = calculate_memory(MODEL, options)
    return [title, gr.update(visible=True, value=pd.DataFrame(data)), gr.update(visible=not has_discussion)]


with gr.Blocks() as demo:
    gr.Markdown("""<h1>Positron Model Universe Explorer</h1>""")
    with gr.Tab("Model Size"):
        gr.Markdown(
            """<h1>Model Memory Calculator</h1>

            This tool will help you calculate how much vRAM is needed to train and perform big model inference
            on a model hosted on the 🤗 Hugging Face Hub. The minimum recommended vRAM needed for a model
            s denoted as the size of the "largest layer", and training of a model is roughly 4x its size (for Adam)."""
        )

        out_text = gr.Markdown()
        mem_out = gr.DataFrame(
            headers=["dtype", "Largest Layer", "Total Size", "Training using Adam"],
            interactive=False,
            visible=True,
        )
        with gr.Row():
            model_name_inp = gr.Textbox(label="Model Name or URL", value="TheBloke/Nous-Hermes-13B-GPTQ")
        with gr.Row():
            library = gr.Radio(["auto", "transformers", "timm"], label="Library", value="auto")
            options = gr.CheckboxGroup(
                ["float32", "float16/bfloat16", "int8", "int4"],
                value="float32",
                label="Model Precision",
            )
            access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
        with gr.Row():
            mem_btn = gr.Button("Calculate Memory Usage")
            post_to_hub = gr.Button(
                value="Report results in this model repo's discussions!\n(Will open in a new tab)", visible=False
            )

        mem_btn.click(
            get_mem_results,
            inputs=[model_name_inp, library, options, access_token],
            outputs=[out_text, mem_out, post_to_hub],
        )

    with gr.Tab("Model Type"):
        gr.Markdown(
            """<h1>Models by Model Task</h1>"""
        )

        with gr.Row():
            task_inp = gr.Dropdown(choices = ["text-generation", "question-answering", "text-classification", "unconditional-image-generation"], 
                value="text-generation", interactive=True, filterable=True, label="Model Task", show_label=True
            )
            stat_inp = gr.Radio(choices = ["downloads", "likes"], value = "downloads", label = "Sorting By", show_label=True, interactive = True)
        
        with gr.Row():
            stats_btn = gr.Button("Pull Stats from HF API")

        stats_btn.click(
            get_model_type_downloads,
            inputs=[task_inp],
            outputs=[stats_out]
        )

        stats_out = gr.Dataframe(row_count=10)
        # gr.Dataframe(get_model_type_likes(), row_count=10)

demo.launch(share=True)