import torch


def mixture_bernoulli_loss(label, log_theta, log_alpha, adj_loss_func,
                           subgraph_idx):
    """
      Compute likelihood for mixture of Bernoulli model

      Args:
        label: E * 1, see comments above
        log_theta: E * D, see comments above
        log_alpha: E * D, see comments above
        adj_loss_func: BCE loss
        subgraph_idx: E * 1, see comments above

      Returns:
        loss: negative log likelihood
    """

    num_subgraph = subgraph_idx.max() + 1
    K = log_theta.shape[1]
    adj_loss = torch.stack(
        [adj_loss_func(log_theta[:, kk], label) for kk in range(K)], dim=1)

    const = torch.zeros(num_subgraph).to(label.device)
    const = const.scatter_add(0, subgraph_idx,
                              torch.ones_like(subgraph_idx).float())

    reduce_adj_loss = torch.zeros(num_subgraph, K).to(label.device)
    reduce_adj_loss = reduce_adj_loss.scatter_add(
        0, subgraph_idx.unsqueeze(1).expand(-1, K), adj_loss)

    reduce_log_alpha = torch.zeros(num_subgraph, K).to(label.device)
    reduce_log_alpha = reduce_log_alpha.scatter_add(
        0, subgraph_idx.unsqueeze(1).expand(-1, K), log_alpha)
    reduce_log_alpha = reduce_log_alpha / const.view(-1, 1)

    log_prob = -reduce_adj_loss + reduce_log_alpha
    log_prob = torch.logsumexp(log_prob, dim=1)
    loss = -log_prob.sum() / float(log_theta.shape[0])

    return loss