import torch


def variational_accuracy(y_predicted, y_target):
    """
    Finds the modal class over the number of samples and evaluates accuracy of that prediction
    Takes logits.

    :param y_predicted: Tensor [examples, samples from the variational posterior, class]
    :param y_target: Tensor [examples]
    :return: float accuracy
    """
    assert len(y_predicted) == len(y_target), "The output y should be the same length as the targets"
    # output of the model is a log_softmax (which is more efficient in the training loop), so we exponentiate
    y_predicted = torch.exp(y_predicted)
    class_prediction = torch.argmax(y_predicted, dim=2)
    modal_prediction, _ = torch.mode(class_prediction, dim=1)
    assert modal_prediction.shape[0] == y_predicted.shape[0], "arg maxes should not have changed the length"
    correct = torch.sum(modal_prediction == y_target)
    assert correct <= len(y_predicted), "Should not be more correct answers than examples."
    return (correct.type(torch.float32) / len(y_predicted)).cpu()