import torch

import torch.nn.functional as F

from .multitask_loss import get_multitask_loss

from .time_aware_loss import get_time_aware_loss

def get_loss(y_pred, y_true, task, time_aware=False):

    if task == "outcome":

        loss = F.binary_cross_entropy(y_pred, y_true[:, 0])

    elif task == "los":

        loss = F.mse_loss(y_pred, y_true[:, 1])

    elif task == "multitask":

        loss = get_multitask_loss(y_pred[:,0], y_pred[:,1], y_true[:,0], y_true[:,1])

    # If use time aware loss:

    if task == "outcome" and time_aware:

        loss = get_time_aware_loss(y_pred, y_true[:, 0], y_true[:, 1])

    return loss