import gradio as gr
from mammal.keys import *

from mammal_demo.demo_framework import MammalObjectBroker    


from mammal_demo.ppi_task import PpiTask
from mammal_demo.dti_task import DtiTask

all_tasks = dict()
all_models= dict()

ppi_task = PpiTask(model_dict = all_models)
all_tasks[ppi_task.name]=ppi_task

tdi_task = DtiTask(model_dict = all_models)
all_tasks[tdi_task.name]=tdi_task

ppi_model = MammalObjectBroker(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=[ppi_task.name])
all_models[ppi_model.name]=ppi_model

tdi_model = MammalObjectBroker(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd", task_list=[tdi_task.name])
all_models[tdi_model.name]=tdi_model


def create_application():
    def task_change(value):
        visibility = [gr.update(visible=(task==value)) for task in all_tasks.keys()]
            # all_tasks[task].demo().visible = 
        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]),*visibility)
        else:
            return (gr.skip,*visibility)
        # return model_name_dropdown
        
       
    with gr.Blocks() as application:
        task_dropdown = gr.Dropdown(choices=["select demo"] + list(all_tasks.keys()))
        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)
        
            



        ppi_demo = all_tasks[ppi_task.name].demo(model_name_widgit = model_name_dropdown)
        # ppi_demo.visible = True
        dtb_demo = all_tasks[tdi_task.name].demo(model_name_widgit = model_name_dropdown)

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


if __name__ == "__main__":
    main()