import torch


def compute_ece(predictions, labels, num_bins=None):
    """ Compute the expected calibration error with binning 

    Args:
        predictions (tensor): a batch of categorical predictions with shape [batch_size, num_classes]
        labels (tensor): a batch of labels with shape [batch_size]
        num_bins (int): number of bins to bin the confidences, set to None for the default choice
        
    Returns:
        tensor: the computed ECE value, a tensor of shape []
    """
    confidence = predictions.max(dim=1)[0]
    correct = (predictions.argmax(dim=1) == labels.to(predictions.device)).type(torch.float32)

    # Sort by confidence 
    ranking = torch.argsort(confidence)
    sorted_confidence = confidence[ranking]
    sorted_correct = correct[ranking]

    if num_bins is None:
        # Heuristically I believe that 1/3 power is a great choice
        # This is because the accuracy of estimating the mean improves sqrt w.r.t. the number of samples
        # The binning error improves linearly w.r.t. the number of bins
        num_bins = int(len(predictions) ** (1./3))  
        
    # Compute the difference between confidence and accuracy
    error = [confidence.mean() - accuracy.mean() for confidence, accuracy in 
             zip(torch.split(sorted_confidence, num_bins), torch.split(sorted_correct, num_bins))]
    error = torch.stack(error)
    
    ece = error.abs().mean()
    return ece