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

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))
tcr_task = all_tasks.register_task(TcrTask(model_dict=all_models))
ps_task = all_tasks.register_task(PsTask(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, tcr_task],
)
all_models.register_model("https://huggingface.co/ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_tox")
all_models.register_model("https://huggingface.co/ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_fda")
all_models.register_model("https://huggingface.co/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.skip, *visibility)
        # return model_name_dropdown

    with gr.Blocks() as application:
        task_dropdown = gr.Dropdown(
            choices=["Select task"] + list(all_tasks.keys()), label="Mammal Task"
        )
        task_dropdown.interactive = True
        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,
        )

        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
            ],
        )

        # def set_demo_vis(main_text):
        #     main_text=main_text
        #     print(f"main text is {main_text}")
        #     return gr.Group(visible=True)
        #     #return gr.Group(visible=(main_text == "PPI"))
        # # , gr.Group(                visible=(main_text == "DTI")            )

        # task_dropdown.change(
        # set_ppi_vis, inputs=task_dropdown, outputs=[ppi_demo]
        # )
        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()