import torch


def evaluate_ece(confidences: torch.Tensor,
                 true_labels: torch.Tensor,
                 n_bins: int = 15) -> float:
    """
    Args:
        confidences (Tensor): a tensor of shape [N, K] of predicted confidences.
        true_labels (Tensor): a tensor of shape [N,] of ground truth labels.
        n_bins (int): the number of bins used by the histrogram binning.

    Returns:
        ece (float): expected calibration error of predictions.
    """
    # predicted labels and its confidences
    pred_confidences, pred_labels = torch.max(confidences, dim=1)

    # fixed binning (n_bins)
    ticks = torch.linspace(0, 1, n_bins + 1)
    bin_lowers = ticks[:-1]
    bin_uppers = ticks[ 1:]

    # compute ECE across bins
    accuracies = pred_labels.eq(true_labels)
    ece = torch.zeros(1, device=confidences.device)
    avg_accuracies = []
    avg_confidences = []
    for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
        in_bin = pred_confidences.gt(
            bin_lower.item()
        ) * pred_confidences.le(
            bin_upper.item()
        )
        prop_in_bin = in_bin.float().mean()
        if prop_in_bin.item() > 0:
            accuracy_in_bin = accuracies[in_bin].float().mean()
            avg_confidence_in_bin = pred_confidences[in_bin].mean()
            ece += torch.abs(
                avg_confidence_in_bin - accuracy_in_bin
            ) * prop_in_bin
            avg_accuracies.append(accuracy_in_bin.item())
            avg_confidences.append(avg_confidence_in_bin.item())
        else:
            avg_accuracies.append(None)
            avg_confidences.append(None)

    return ece.item()