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


def tasks_and_models():
    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.dti_bindingdb_pkd_peer",
    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"
)
    
    return all_tasks,all_models