import torch


def get_TFPN_dict(preds, labels, true_label=1, as_float=False):
    """
    Given predictions and labels, returns a dictionary with TP, TN, FP and FN for a given class label.
    """
    eval_dict = dict()
    eval_dict["TP"] = torch.logical_and(preds == true_label, preds == labels)
    eval_dict["TN"] = torch.logical_and(preds != true_label, preds == labels)
    eval_dict["FP"] = torch.logical_and(preds == true_label, preds != labels)
    eval_dict["FN"] = torch.logical_and(preds != true_label, preds != labels)
    eval_dict = {key: eval_dict[key].long().sum() for key in eval_dict}
    if as_float:
        eval_dict = {key: eval_dict[key].float() for key in eval_dict}
    return eval_dict