import torch


def compute_score(logits, labels):
    """Computes score of the prediction.

    Ground truth soft scores for labels is calculated as per:
    https://github.com/hengyuan-hu/bottom-up-attention-vqa/blob/master/tools/compute_softscore.py#L80
    and http://visualqa.org/evaluation.html
    Args:
        logits:
            Tensor containing predicted score for each answer candidiate
            predicted. Scores are pre-sigmoid and not between 0 and 1.
            Tensor is of dimension (batch size, answer vocabulary size).
        labels:
            Ground truth soft scores for each answer in the set of candidiate
            answers.
            Soft scores/ accuracies are between [0, 1].
            Tensor is of dimension (batch size, answer vocabulary size).
    """
    # logits has dimension [batch_size, num_ans_candidiates]
    # Each sample in the batch is of dimension num_ans_candidiates and each
    # element is the predicted score for the corresponding answer.

    # Finding the element index (offset) with max score for each sample in the
    # batch (i.e. argmax).
    # logits will now have dimension [batch_size] i.e. batch_size * 1.
    logits = torch.max(logits, 1)[1].data

    # Making zero vector of dimension [batch_size, num_ans_candidiates].
    one_hots = torch.zeros(*labels.size()).cuda()

    # Making a one-hot vector of dimension [batch_size, num_ans_candidiates].
    # For each sample the value at an index = index with max score, will be
    # set to 1.
    # This one hot vector is essentially encodes the most probable answer in
    # the vocabulary as predicted by the model.
    one_hots.scatter_(1, logits.view(-1, 1), 1)

    # Score will be a vector of dimension [batch_size, num_ans_candidiates].
    # Each sample will be a vector containing all zeroes, except at the index
    # of the predicted answer which contains the soft score/accuracy of that
    # answer calculated as per http://visualqa.org/evaluation.html.
    scores = one_hots * labels
    return scores