import torch


def jaccard_index(pred, target, n_classes):
    """
    Compute jaccard index (intersection over union) metric

    :param pred: torch.Tensor, prediction mask
    :param target: torch.Tensor, target mask
    :param n_classes: int, number of classes
    :return: torch.Tensor, jaccard index per class
    """
    ious = []
    pred = pred.view(-1)
    target = target.view(-1)
    assert pred.shape == target.shape, '{} != {}'.format(pred.shape, target.shape)

    # Ignore IoU for background class ("0")
    # for cls in range(1, n_classes):  # This goes from 1:n_classes-1 -> class "0" is ignored
    for cls in range(n_classes):  # This goes from 0:n_classes-1 -> class "0" is NOT ignored
        pred_inds = pred == cls
        target_inds = target == cls
        intersection = (pred_inds[target_inds]).long().sum().item()
        union = pred_inds.long().sum().item() + target_inds.long().sum().item() - intersection
        if union == 0:
            ious.append(float('nan'))  # If there is no ground truth, do not include in evaluation
        else:
            ious.append(float(intersection) / float(max(union, 1)))
    return torch.tensor(ious)