import logging
from celery.signals import worker_process_init
from celery import shared_task
from epdb.models import Pathway, Node, Edge, EPModel, Setting
from epdb.logic import SPathway

from utilities.chem import FormatConverter

logger = logging.getLogger(__name__)


@shared_task(queue='background')
def mul(a, b):
    return a * b


@shared_task(queue='predict')
def predict_simple(model_pk: int, smiles: str):
    mod = EPModel.objects.get(id=model_pk)
    res = mod.predict(smiles)
    return res


@shared_task(queue='background')
def send_registration_mail(user_pk: int):
    pass


@shared_task(queue='model')
def build_model(model_pk: int):
    mod = EPModel.objects.get(id=model_pk)
    X, y = mod.build_dataset()
    mod.build_model(X, y)


@shared_task(queue='model')
def evaluate_model(model_pk: int):
    mod = EPModel.objects.get(id=model_pk)
    mod.evaluate_model()


@shared_task(queue='predict')
def predict(pw_pk: int, pred_setting_pk: int):
    pw = Pathway.objects.get(id=pw_pk)
    setting = Setting.objects.get(id=pred_setting_pk)
    spw = SPathway(prediction_setting=setting, persist=pw)
    level = 0
    while not spw.done:
        spw.predict_step(from_depth=level)
        level += 1