import torch


def get_dice(predicted, labels, n_class, area=None):
    """
    Dice socre.

    Parameters
    __________
    predicted : torch.tensor, shape = [batch, n_vertex]
        Model outputs.
    labels : torch.tensor, shape = [batch, n_vertex]
        True labels.
    n_class : int
        # of classes.
    area : torch.tensor, shape = [n_vertex]
        Vertex-wise area.

    Returns
    _______
    dice : torch.tensor, shape = [batch, n_class]
        Batch-wise Dice score.
    """

    if area is None:
        area = 1
    num_batch = predicted.shape[0]
    batch_numer = torch.zeros(num_batch, n_class)
    batch_denom = torch.zeros(num_batch, n_class)

    for i in range(n_class):
        batch_numer[:, i] = torch.mul(area, ((predicted == i) & (labels == i)).int()).sum(dim=1)
        batch_denom[:, i] = torch.mul(area, (labels == i).int()).sum(dim=1) + torch.mul(
            area, (predicted == i).int()
        ).sum(dim=1)

    return 2 * batch_numer / batch_denom