# import torch
# from sklearn.metrics import roc_auc_score, roc_curve, average_precision_score
#
# def roc_auc_compute_fn(y_preds: torch.Tensor, y_targets: torch.Tensor):
#     y_true = y_targets.cpu().numpy()
#     y_pred = y_preds.cpu().numpy()
#     return roc_auc_score(y_true, y_pred)
#
#
# def roc_auc_curve_compute_fn(y_preds: torch.Tensor, y_targets: torch.Tensor):
#     y_true = y_targets.cpu().numpy()
#     y_pred = y_preds.cpu().numpy()
#     return roc_curve(y_true, y_pred)
#
# def precision_recall_curve_compute_fn(y_preds: torch.Tensor, y_targets: torch.Tensor):
#     try:
#         from sklearn.metrics import precision_recall_curve
#     except ImportError:
#         raise ModuleNotFoundError("This contrib module requires scikit-learn to be installed.")
#
#     y_true = y_targets.cpu().numpy()
#     y_pred = y_preds.cpu().numpy()
#     return precision_recall_curve(y_true, y_pred)
#
# def average_precision_compute_fn(y_preds: torch.Tensor, y_targets: torch.Tensor) -> float:
#     y_true = y_targets.cpu().numpy()
#     y_pred = y_preds.cpu().numpy()
#     return average_precision_score(y_true, y_pred)
#
#
# class MultiLabelConfusionMatrix():
#     """Calculates a confusion matrix for multi-labelled, multi-class data."""
#     def __init__(self,num_classes,output_transform,device=torch.device("cpu"),normalized = False,):
#         if num_classes <= 1:
#             raise ValueError("Argument num_classes needs to be > 1")
#         self.num_classes = num_classes
#         self.normalized = normalized
#
#     def reset(self):
#         self.confusion_matrix = torch.zeros(self.num_classes, 2, 2, dtype=torch.int64, device=self._device)
#
#     def update(self, output):
#         self._check_input(output)
#         y_pred, y = output[0].detach(), output[1].detach()
#
#         y_reshaped = y.transpose(0, 1).reshape(self.num_classes, -1)
#         y_pred_reshaped = y_pred.transpose(0, 1).reshape(self.num_classes, -1)
#
#         y_total = y_reshaped.sum(dim=1)
#         y_pred_total = y_pred_reshaped.sum(dim=1)
#
#         tp = (y_reshaped * y_pred_reshaped).sum(dim=1)
#         fp = y_pred_total - tp
#         fn = y_total - tp
#         tn = y_reshaped.shape[1] - tp - fp - fn
#
#         self.confusion_matrix += torch.stack([tn, fp, fn, tp], dim=1).reshape(-1, 2, 2).to(self._device)
#
#     def compute(self):
#         if self.normalized:
#             conf = self.confusion_matrix.to(dtype=torch.float64)
#             sums = conf.sum(dim=(1, 2))
#             return conf / sums[:, None, None]
#
#         return self.confusion_matrix
#
#     def _check_input(self, output):
#         y_pred, y = output[0].detach(), output[1].detach()
#
#         if y_pred.ndimension() < 2:
#             raise ValueError(
#                 f"y_pred must at least have shape (batch_size, num_classes (currently set to {self.num_classes}), ...)"
#             )
#
#         if y.ndimension() < 2:
#             raise ValueError(
#                 f"y must at least have shape (batch_size, num_classes (currently set to {self.num_classes}), ...)"
#             )
#
#         if y_pred.shape[0] != y.shape[0]:
#             raise ValueError(f"y_pred and y have different batch size: {y_pred.shape[0]} vs {y.shape[0]}")
#
#         if y_pred.shape[1] != self.num_classes:
#             raise ValueError(f"y_pred does not have correct number of classes: {y_pred.shape[1]} vs {self.num_classes}")
#
#         if y.shape[1] != self.num_classes:
#             raise ValueError(f"y does not have correct number of classes: {y.shape[1]} vs {self.num_classes}")
#
#         if y.shape != y_pred.shape:
#             raise ValueError("y and y_pred shapes must match.")
#
#         valid_types = (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
#         if y_pred.dtype not in valid_types:
#             raise ValueError(f"y_pred must be of any type: {valid_types}")
#
#         if y.dtype not in valid_types:
#             raise ValueError(f"y must be of any type: {valid_types}")
#
#         if not torch.equal(y_pred, y_pred**2):
#             raise ValueError("y_pred must be a binary tensor")
#
#         if not torch.equal(y, y**2):
#             raise ValueError("y must be a binary tensor")