import torch


def fast_compute_class_metric(test_preds, test_sols, class_metrics = ('accuracy', 'recall', 'precision')):
    """
    Computes manually (i.e. without sklearn functions) accuracy, recall and predicision.

    Args:
        test_preds: numpy array/ torch tensor of size N with discrete output vals
        test_sols: numpy array/torch tensor of size N with binary labels
        class_metrics: tuple with a subset of values from ('accuracy', 'recall', 'precision') indicating which
        metrics to report

    Returns:
        dictionary with values of accuracy, recall and precision
    """
    with torch.no_grad():

        TP = ((test_sols == 1) & (test_preds == 1)).sum().float()
        FP = ((test_sols == 0) & (test_preds == 1)).sum().float()
        TN = ((test_sols == 0) & (test_preds == 0)).sum().float()
        FN = ((test_sols == 1) & (test_preds == 0)).sum().float()

        accuracy = (TP + TN) / (TP + FP + TN + FN)
        recall = TP / (TP + FN) if TP + FN > 0 else torch.tensor(0)
        precision = TP / (TP + FP) if TP + FP > 0 else torch.tensor(0)

    class_metrics_dict =  {'accuracy': accuracy.item(), 'recall': recall.item(), 'precision': precision.item()}
    class_metrics_dict = {met_name: class_metrics_dict[met_name] for met_name in class_metrics}

    return class_metrics_dict