import gradio as gr

from mammal_demo.demo_framework import (
    ModelRegistry,
    TaskRegistry,
)
from mammal_demo.dti_task import DtiTask
from mammal_demo.ppi_task import PpiTask
from mammal_demo.ps_task import PsTask
from mammal_demo.tcr_task import TcrTask

MAIN_MARKDOWN_TEXT = """

The **[ibm/biomed.omics.bl.sm.ma-ted-458m](https://huggingface.co/models?sort=trending&search=ibm%2Fbiomed.omics.bl)**  model family is a biomedical foundation model and its finetuned variants trained on over 2 billion biological samples across multiple modalities, including proteins, small molecules, and single-cell gene data.  
Designed for robust performance, it achieves state-of-the-art results over a variety of tasks across the entire drug discovery pipeline and the diverse biomedical domains.  

Based on the **M**olecular **A**ligned **M**ulti-**M**odal **A**rchitecture and **L**anguage (**MAMMAL**), a flexible, multi-domain architecture with an adaptable task prompt syntax.  
The syntax allows for dynamic combinations of tokens and scalars, enabling classification, regression, and generation tasks either within a single domain or with cross-domain entities.

This page demonstraits a variety of drug discovery and biomedical tasks for the model family.  Select the task to access the specific demos.
"""


all_tasks = TaskRegistry()
all_models = ModelRegistry()

# first create the required tasks
# Note that the tasks need access to the models, as the model to use depends on the state of the widget
# we pass the all_models dict and update it when we actualy have the models.

ppi_task = all_tasks.register_task(PpiTask(model_dict=all_models))
tdi_task = all_tasks.register_task(DtiTask(model_dict=all_models))
ps_task = all_tasks.register_task(PsTask(model_dict=all_models))
tcr_task = all_tasks.register_task(TcrTask(model_dict=all_models))

# create the model holders. hold the model and the tokenizer, lazy download
# note that the list of relevent tasks needs to be stated.
all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd",
    task_list=[tdi_task],
)
all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.tcr_epitope_bind",
    task_list=[tcr_task],
)
all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.protein_solubility",
    task_list=[ps_task],
)
all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m",
    task_list=[ppi_task],
)
all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_tox"
)
all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_fda"
)
all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_bbbp"
)


def create_application():
    def task_change(value):
        visibility = [gr.update(visible=(task == value)) for task in all_tasks.keys()]
        choices = [
            model_name
            for model_name, model in all_models.items()
            if value in model.tasks
        ]
        if choices:
            return (
                gr.update(choices=choices, value=choices[0], visible=True),
                *visibility,
            )
        else:
            return (gr.update(visible=False, value=None), *visibility)

    def model_change(value):
        return gr.update(
            value=f'[<span style="font-size:4em;">🤗</span>to model](https://huggingface.co/{value})',
            visible=value is not None,
        )

    with gr.Blocks(theme="Zarkel/IBM_Carbon_Theme") as application:
        application_main_markdown=gr.Markdown(MAIN_MARKDOWN_TEXT, visible=True)
        task_dropdown = gr.Dropdown(
            choices=["Select task"] + list(all_tasks.keys()),
            label="Mammal Task",
        )
        task_dropdown.interactive = True
        with gr.Row():
            model_name_dropdown = gr.Dropdown(
                choices=[
                    model_name
                    for model_name, model in all_models.items()
                    if task_dropdown.value in model.tasks
                ],
                interactive=True,
                label="Matching Mammal models",
                visible=False,
                scale=10,
            )
            goto_card_button = gr.Markdown(
                "Link to model card",
                visible=False,
            )

            def echo(value):
                print(value)
                return value

            # goto_card_button.click(
            #     fn=None,
            #     inputs=model_name_dropdown,
            #     js=f"(model_name_dropdown) => {{ window.open('https://huggingface.co/{model_name_dropdown}', '_blank') }}",
            #     )

            model_name_dropdown.change(
                model_change, inputs=[model_name_dropdown], outputs=[goto_card_button]
            )

        task_dropdown.change(
            task_change,
            inputs=[task_dropdown],
            outputs=[model_name_dropdown]
            + [
                all_tasks[task].demo(model_name_widgit=model_name_dropdown)
                for task in all_tasks
            ],
        )

        return application


full_demo = None


def main():
    global full_demo
    full_demo = create_application()
    full_demo.launch(show_error=True, share=False)
    # full_demo.launch(show_error=True, share=True)


if __name__ == "__main__":
    main()