## Define the loss function
import torch.nn.functional as F

from .distribution import NegativeBinomial, ZeroInflatedNegativeBinomial


def rec_loss(x_hat,x,logtheta,zi_logit,distribution):
    if distribution == 'nb':
        loss  = -NegativeBinomial(mu=x_hat, theta=logtheta.exp()).log_prob(x).sum(-1).mean()
    elif distribution == 'zinb':
        loss  = -ZeroInflatedNegativeBinomial(mu=x_hat, theta=logtheta.exp(),zi_logits=zi_logit).log_prob(x).sum(-1).mean()
    else:
        loss = F.mse_loss(x_hat, x)

    return loss

def ce_loss(pred_label, true_label):
    return F.cross_entropy(pred_label, true_label)
