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
from mammal_demo.molnet_task import MolnetTask

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_name = all_tasks.register_task(PpiTask(model_dict=all_models))
    tdi_task_name = all_tasks.register_task(DtiTask(model_dict=all_models))
    ps_task_name = all_tasks.register_task(PsTask(model_dict=all_models))
    tcr_task_name = all_tasks.register_task(TcrTask(model_dict=all_models))
    bbbp_task =  MolnetTask(model_dict=all_models,task_name="BBBP", name= "Blood-Brain Barrier Penetration")
    bbbp_task.markup_text = """
# Mammal based small molecule blood-brain barrier penetration demonstration

Given a drug (in SMILES), estimate the likelihood that it will penetrate the Blood-Brain Barrier.
"""
    bbbp_task_name = all_tasks.register_task(bbbp_task)
    
    toxicity_task = MolnetTask(model_dict=all_models,task_name="TOXICITY", name= "Drug Toxicity Trials Failer")
    toxicity_task.markup_text = """
# Mammal based small molecule toxicity trials failer estimation demonstration

Given a drug (in SMILES), estimate the likelihood that it will fail in clinical toxicity trials.
"""
    toxicity_task_name = all_tasks.register_task(toxicity_task)
    
    
    
    fda_appr_task=MolnetTask(model_dict=all_models,task_name="FDA_APPR", name="drug FDA approval demonstration")
    fda_appr_task.markup_text = """
# Mammal based small molecule drug FDA approval demonstration

Given a drug (in SMILES), estimate the likelihood that it will be approved by the FDA.
    """
    fda_appr_task_name = all_tasks.register_task(fda_appr_task)
   

# 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_name],
)
    all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd_peer",
    task_list=[tdi_task_name],
)

    all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.tcr_epitope_bind",
    task_list=[tcr_task_name],
)
    all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.protein_solubility",
    task_list=[ps_task_name],
)
    all_models.register_model(
    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m",
    task_list=[ppi_task_name],
)
    all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_tox",
    task_list=[toxicity_task_name]
)
    all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_fda",
    task_list=[fda_appr_task_name]
)
    all_models.register_model(
    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_bbbp",
    task_list=[bbbp_task_name],
)
    
    return all_tasks,all_models