import gradio as gr

import mammal_demo
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 [**MAMMAL** - **M**olecular **A**ligned **M**ulti-**M**odal **A**rchitecture and **L**anguage](https://arxiv.org/abs/2410.22367v2), 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, all_models = mammal_demo.tasks_and_models()


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:
            active = len(choices)>1
            return (
                gr.update(choices=choices, value=choices[0], interactive=active, visible=True, label=f"Matching Mammal models ({len(choices)})",),
                *visibility,
            )
        else:
            return (gr.update(visible=False, value=None, label="No Matching Mammal models"), *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="matanninio/IBM_Carbon_Theme@0.0.5") as application:
        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_selector = gr.Radio(
                choices=[
                    model_name
                    for model_name, model in all_models.items()
                    if task_dropdown.value in model.tasks
                ],
                interactive=True,
                label="",
                visible=False,
                scale=10,
            )
            goto_card_button = gr.Markdown(
                "Link to model card",
                visible=False,
            )

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

        task_dropdown.change(
            task_change,
            inputs=[task_dropdown],
            outputs=[model_name_selector]
            + [
                all_tasks[task].demo(model_name_widgit=model_name_selector)
                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()