import torch


def create_class_histograms(pred_prob, target, num_bins):
    """
    Creates two histograms to contain total positive or true samples
    for a given class at a given score value and total samples at a
    given predicted probability.  Then for a given threshold we can compute
    precision = true_samples_above_threshold /
    total_samples_above_threshold and recall =
    true_samples_above_threshold / true_samples. Output is expected to
    be prediction probabilities.

    The pred_prob vector is num_classes x batchsize_per_replica, the target vector is
    num_classes x 1, should contain integers corresponding to each class.
    Output is two tensors of same size: num_bins x num_classes
    where the bins are equispaced from 0 to 1 across the score.
    The class_hist[:, c] contains true example histogram for class c
    prediction probabilities while total_hist[:, c] contains histogram
    for all samples for class c prediction probabilities.
    """

    # assertions:
    assert torch.is_tensor(pred_prob)
    assert torch.is_tensor(target)
    assert pred_prob.size()[0] == target.size()[0], "%s pred_prob, %s target" % (
        str(pred_prob.size()),
        str(target.size()),
    )
    assert (
        pred_prob.lt(0.0).sum() == 0 and pred_prob.gt(1.0).sum() == 0
    ), "Prediction probability must be between 0 and 1"

    num_classes = pred_prob.size()[1]
    class_hist = torch.zeros([num_bins, num_classes], dtype=torch.int64)
    total_hist = torch.zeros([num_bins, num_classes], dtype=torch.int64)
    for c in range(num_classes):
        total_hist[:, c] = torch.histc(
            pred_prob[:, c], bins=num_bins, min=0.0, max=1.0
        ).long()
        class_hist[:, c] = torch.histc(
            pred_prob[:, c][(target == c).nonzero().squeeze(1)],
            bins=num_bins,
            min=0.0,
            max=1.0,
        ).long()

    return class_hist, total_hist