Traceback (most recent call last):
  File "e:\LLMS\Fine-tuning\RLPE\metrics\Imageclassification.py", line 582, in <module>       
    metrics_binary = metrics_calc_binary.get_all_metrics()
  File "e:\LLMS\Fine-tuning\RLPE\metrics\Imageclassification.py", line 562, in get_all_metrics
    'Matthews Correlation Coefficient': self.matthews_correlation_coefficient(),
  File "e:\LLMS\Fine-tuning\RLPE\metrics\Imageclassification.py", line 396, in matthews_correlation_coefficient
    denominator = torch.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn) + 1e-12)
TypeError: sqrt(): argument 'input' (position 1) must be Tensor, not float

import torch
from torch import Tensor
from typing import Optional, Dict, Any
import warnings


class MetricsCalculator:
    """
    A class to compute various image classification metrics using PyTorch.

    Attributes:
        preds (Tensor): Model predictions (logits or probabilities).
        labels (Tensor): Ground truth labels.
        num_classes (Optional[int]): Number of classes. Required for multi-class metrics.
    """

    def __init__(self, preds: Tensor, labels: Tensor, num_classes: Optional[int] = None):
        """
        Initializes the MetricsCalculator with predictions and labels.

        Args:
            preds (Tensor): Model predictions (logits or probabilities).
            labels (Tensor): Ground truth labels.
            num_classes (Optional[int], optional): Number of classes. Defaults to None.

        Raises:
            ValueError: If preds and labels have mismatched shapes or invalid dimensions.
        """
        if preds.ndim != labels.ndim and labels.ndim not in [1, 2]:
            raise ValueError(
                "Predictions and labels must have the same number of dimensions."
            )
        if preds.size(0) != labels.size(0):
            raise ValueError("Predictions and labels must have the same number of samples.")

        self.preds = preds
        self.labels = labels
        self.num_classes = num_classes

        if self.num_classes is None:
            if self.preds.ndim > 1:
                self.num_classes = self.preds.size(1)
            else:
                self.num_classes = int(torch.max(labels).item()) + 1

    def _to_probabilities(self) -> Tensor:
        """
        Converts logits to probabilities if necessary.

        Returns:
            Tensor: Probability distributions.
        """
        if self.preds.ndim > 1:
            return torch.softmax(self.preds, dim=1)
        else:
            # Assuming binary classification with logits
            return torch.sigmoid(self.preds)

    def _compute_true_pred_positive(self):
        """
        Computes true positives and predicted positives for each class.

        Returns:
            Tuple[Tensor, Tensor]: true_positive and predicted_positive tensors.
        """
        with torch.no_grad():
            if self.preds.ndim > 1:
                predicted_labels = torch.argmax(self.preds, dim=1)
            else:
                predicted_labels = (torch.sigmoid(self.preds) >= 0.5).long()

            true_positive = torch.zeros(self.num_classes, device=self.labels.device)
            predicted_positive = torch.zeros(self.num_classes, device=self.labels.device)

            for cls in range(self.num_classes):
                true_positive[cls] = ((predicted_labels == cls) & (self.labels == cls)).sum()
                predicted_positive[cls] = (predicted_labels == cls).sum()

            return true_positive, predicted_positive

    def _compute_true_actual_positive(self):
        """
        Computes true positives and actual positives for each class.

        Returns:
            Tuple[Tensor, Tensor]: true_positive and actual_positive tensors.
        """
        with torch.no_grad():
            if self.preds.ndim > 1:
                predicted_labels = torch.argmax(self.preds, dim=1)
            else:
                predicted_labels = (torch.sigmoid(self.preds) >= 0.5).long()

            true_positive = torch.zeros(self.num_classes, device=self.labels.device)
            actual_positive = torch.zeros(self.num_classes, device=self.labels.device)

            for cls in range(self.num_classes):
                true_positive[cls] = ((predicted_labels == cls) & (self.labels == cls)).sum()
                actual_positive[cls] = (self.labels == cls).sum()

            return true_positive, actual_positive

    def accuracy(self) -> float:
        """
        Computes the accuracy of predictions.

        Returns:
            float: Accuracy score.
        """
        with torch.no_grad():
            if self.preds.ndim > 1:
                predicted_labels = torch.argmax(self.preds, dim=1)
            else:
                predicted_labels = (torch.sigmoid(self.preds) >= 0.5).long()
            correct = (predicted_labels == self.labels).sum().item()
            total = self.labels.size(0)
            return correct / total if total > 0 else 0.0

    def precision(self, average: Optional[str] = 'macro') -> float:
        """
        Computes the precision of predictions.

        Args:
            average (str, optional): Averaging method ('macro', 'micro', 'weighted'). Defaults to 'macro'.
                                     If None, returns precision per class.

        Returns:
            float or Tensor: Precision score or tensor of precision per class.

        Raises:
            ValueError: If an unsupported averaging method is provided.
        """
        true_positive, predicted_positive = self._compute_true_pred_positive()
        with torch.no_grad():
            precision = true_positive / (predicted_positive + 1e-12)
            if average is None:
                return precision
            elif average == 'macro':
                return precision.mean().item()
            elif average == 'micro':
                return true_positive.sum().item() / (predicted_positive.sum().item() + 1e-12)
            elif average == 'weighted':
                support = (self.labels.unsqueeze(1) == torch.arange(self.num_classes).to(self.labels.device)).sum(dim=0)
                return (precision * support).sum().item() / (support.sum().item() + 1e-12)
            else:
                raise ValueError("Unsupported average type. Choose from 'macro', 'micro', 'weighted'.")

    def recall(self, average: Optional[str] = 'macro') -> float:
        """
        Computes the recall of predictions.

        Args:
            average (str, optional): Averaging method ('macro', 'micro', 'weighted'). Defaults to 'macro'.
                                     If None, returns recall per class.

        Returns:
            float or Tensor: Recall score or tensor of recall per class.

        Raises:
            ValueError: If an unsupported averaging method is provided.
        """
        true_positive, actual_positive = self._compute_true_actual_positive()
        with torch.no_grad():
            recall = true_positive / (actual_positive + 1e-12)
            if average is None:
                return recall
            elif average == 'macro':
                return recall.mean().item()
            elif average == 'micro':
                return true_positive.sum().item() / (actual_positive.sum().item() + 1e-12)
            elif average == 'weighted':
                support = (self.labels.unsqueeze(1) == torch.arange(self.num_classes).to(self.labels.device)).sum(dim=0)
                return (recall * support).sum().item() / (support.sum().item() + 1e-12)
            else:
                raise ValueError("Unsupported average type. Choose from 'macro', 'micro', 'weighted'.")

    def f1_score(self, average: str = 'macro') -> float:
        """
        Computes the F1 score of predictions.

        Args:
            average (str, optional): Averaging method ('macro', 'micro', 'weighted'). Defaults to 'macro'.

        Returns:
            float: F1 score.

        Raises:
            ValueError: If an unsupported averaging method is provided.
        """
        precision = self.precision(average=None)
        recall = self.recall(average=None)
        with torch.no_grad():
            f1 = 2 * (precision * recall) / (precision + recall + 1e-12)
            if average == 'macro':
                return f1.mean().item()
            elif average == 'micro':
                return (2 * precision.sum() * recall.sum()) / (precision.sum() + recall.sum() + 1e-12)
            elif average == 'weighted':
                support = (self.labels.unsqueeze(1) == torch.arange(self.num_classes).to(self.labels.device)).sum(dim=0)
                return (f1 * support).sum().item() / (support.sum().item() + 1e-12)
            else:
                raise ValueError("Unsupported average type. Choose from 'macro', 'micro', 'weighted'.")

    def confusion_matrix(self) -> Tensor:
        """
        Computes the confusion matrix.

        Returns:
            Tensor: Confusion matrix of shape (num_classes, num_classes).
        """
        with torch.no_grad():
            if self.preds.ndim > 1:
                predicted_labels = torch.argmax(self.preds, dim=1)
            else:
                predicted_labels = (torch.sigmoid(self.preds) >= 0.5).long()

            cm = torch.zeros((self.num_classes, self.num_classes), dtype=torch.int64, device=self.labels.device)
            for t, p in zip(self.labels.view(-1), predicted_labels.view(-1)):
                if 0 <= t < self.num_classes and 0 <= p < self.num_classes:
                    cm[t, p] += 1
            return cm

    def roc_auc(self) -> float:
        """
        Computes the ROC-AUC score for binary classification.

        Returns:
            float: ROC-AUC score.

        Raises:
            ValueError: If not a binary classification problem.
        """
        if self.num_classes != 2:
            raise ValueError("ROC-AUC is only implemented for binary classification.")

        with torch.no_grad():
            probs = self._to_probabilities().squeeze()
            labels = self.labels.float()
            # Sort by predicted probabilities
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            sorted_labels = labels[sorted_indices]
            # Compute True Positive Rate and False Positive Rate
            tp = torch.cumsum(sorted_labels, dim=0)
            fp = torch.cumsum(1 - sorted_labels, dim=0)
            tpr = tp / (labels.sum() + 1e-12)
            fpr = fp / ((1 - labels).sum() + 1e-12)
            # Compute AUC using the trapezoidal rule
            auc = torch.trapz(tpr, fpr).item()
            return auc

    def pr_auc(self) -> float:
        """
        Computes the Precision-Recall AUC for binary classification.

        Returns:
            float: PR-AUC score.

        Raises:
            ValueError: If not a binary classification problem.
        """
        if self.num_classes != 2:
            raise ValueError("PR-AUC is only implemented for binary classification.")

        with torch.no_grad():
            probs = self._to_probabilities().squeeze()
            labels = self.labels.float()
            # Sort by predicted probabilities
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            sorted_labels = labels[sorted_indices]
            tp = torch.cumsum(sorted_labels, dim=0)
            fp = torch.cumsum(1 - sorted_labels, dim=0)
            precision = tp / (tp + fp + 1e-12)
            recall = tp / (labels.sum() + 1e-12)
            # Compute AUC using the trapezoidal rule
            # To compute PR-AUC, integrate precision w.r.t. recall
            # Since recall is sorted, we can use torch.trapz
            auc = torch.trapz(precision, recall).item()
            return auc

    def top_k_accuracy(self, k: int = 5) -> float:
        """
        Computes the Top-k accuracy.

        Args:
            k (int, optional): The 'k' in Top-k. Defaults to 5.

        Returns:
            float: Top-k accuracy score.

        Raises:
            ValueError: If k is invalid.
        """
        if k <= 0:
            raise ValueError("k must be a positive integer.")
        if self.preds.ndim == 1 and k > 1:
            raise ValueError("Top-k accuracy requires multi-class predictions.")

        with torch.no_grad():
            if self.preds.ndim > 1:
                topk = torch.topk(self.preds, k, dim=1).indices
                correct = topk.eq(self.labels.view(-1, 1).expand_as(topk))
                correct_total = correct.any(dim=1).sum().item()
            else:
                probs = torch.sigmoid(self.preds)
                topk = (probs >= torch.topk(probs, k)[0][-1]).long()
                correct_total = (topk.squeeze() == self.labels).sum().item()
            total = self.labels.size(0)
            return correct_total / total if total > 0 else 0.0

    def log_loss(self) -> float:
        """
        Computes the Logarithmic Loss.

        Returns:
            float: Log loss score.
        """
        with torch.no_grad():
            probs = self._to_probabilities()
            if self.num_classes == 2 and self.preds.ndim == 1:
                probs = probs.clamp(min=1e-12, max=1 - 1e-12)
                loss = -(self.labels.float() * torch.log(probs) + (1 - self.labels.float()) * torch.log(1 - probs))
                return loss.mean().item()
            else:
                one_hot = torch.nn.functional.one_hot(self.labels, num_classes=self.num_classes).float()
                probs = probs.clamp(min=1e-12, max=1 - 1e-12)
                loss = -torch.sum(one_hot * torch.log(probs), dim=1)
                return loss.mean().item()

    def cross_entropy_loss(self) -> float:
        """
        Computes the Cross-Entropy Loss.

        Returns:
            float: Cross-entropy loss score.
        """
        with torch.no_grad():
            if self.preds.ndim == 1 and self.num_classes == 2:
                # Binary classification
                loss = torch.nn.functional.binary_cross_entropy_with_logits(self.preds, self.labels.float())
            else:
                # Multi-class classification
                loss = torch.nn.functional.cross_entropy(self.preds, self.labels)
            return loss.item()

    def mean_squared_error(self) -> float:
        """
        Computes the Mean Squared Error (MSE).

        Returns:
            float: MSE score.
        """
        with torch.no_grad():
            if self.preds.ndim == 1:
                preds = torch.sigmoid(self.preds)
            elif self.preds.ndim > 1:
                preds = torch.softmax(self.preds, dim=1)
                one_hot = torch.nn.functional.one_hot(self.labels, num_classes=self.num_classes).float()
                preds = preds[range(self.labels.size(0)), self.labels]
            else:
                preds = self.preds
            targets = self.labels.float()
            mse = torch.mean((preds - targets) ** 2).item()
            return mse

    def kappa_score(self) -> float:
        """
        Computes Cohen’s Kappa statistic.

        Returns:
            float: Kappa score.
        """
        with torch.no_grad():
            cm = self.confusion_matrix().float()
            total = cm.sum()
            sum0 = cm.sum(dim=1)
            sum1 = cm.sum(dim=0)
            expected = torch.outer(sum0, sum1) / total
            kappa = (cm.diag().sum() - expected.diag().sum()) / (total - expected.diag().sum() + 1e-12)
            return kappa.item()

    def matthews_correlation_coefficient(self) -> float:
        """
        Computes the Matthews Correlation Coefficient (MCC).

        Returns:
            float: MCC score.
        """
        cm = self.confusion_matrix()
        with torch.no_grad():
            if self.num_classes == 2:
                tp = cm[1, 1].item()
                tn = cm[0, 0].item()
                fp = cm[0, 1].item()
                fn = cm[1, 0].item()
                numerator = tp * tn - fp * fn
                denominator = torch.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn) + 1e-12)
                return numerator / denominator if denominator != 0 else 0.0
            else:
                # For multi-class, use generalized MCC
                t_sum = cm.sum(dim=1)
                p_sum = cm.sum(dim=0)
                c = cm.diag().sum()
                s = cm.sum()
                numerator = (c * s) - torch.dot(t_sum, p_sum)
                denominator = torch.sqrt((s ** 2 - torch.dot(p_sum, p_sum)) * (s ** 2 - torch.dot(t_sum, t_sum)) + 1e-12)
                return (numerator / denominator).item() if denominator != 0 else 0.0

    def specificity(self) -> float:
        """
        Computes the specificity (True Negative Rate).

        Returns:
            float: Specificity score.

        Raises:
            ValueError: If not a binary classification problem.
        """
        if self.num_classes != 2:
            raise ValueError("Specificity is only implemented for binary classification.")

        with torch.no_grad():
            cm = self.confusion_matrix()
            tn = cm[0, 0].item()
            fp = cm[0, 1].item()
            specificity = tn / (tn + fp + 1e-12)
            return specificity

    def jaccard_index(self) -> float:
        """
        Computes the Jaccard Index (IoU).

        Returns:
            float: Jaccard Index score.
        """
        with torch.no_grad():
            if self.num_classes == 2 and self.preds.ndim == 1:
                preds = (torch.sigmoid(self.preds) >= 0.5).long()
                intersection = (preds & self.labels).sum().item()
                union = (preds | self.labels).sum().item()
                return intersection / union if union > 0 else 0.0
            else:
                # Multi-class IoU
                cm = self.confusion_matrix()
                intersection = cm.diag()
                union = cm.sum(dim=1) + cm.sum(dim=0) - cm.diag()
                iou = (intersection / (union + 1e-12)).mean().item()
                return iou

    def balanced_accuracy(self) -> float:
        """
        Computes the Balanced Accuracy.

        Returns:
            float: Balanced accuracy score.
        """
        recall = self.recall(average='macro')
        return recall

    def true_negative_rate(self) -> float:
        """
        Computes the True Negative Rate (TNR), equivalent to specificity.

        Returns:
            float: TNR score.
        """
        return self.specificity()

    def false_positive_rate(self) -> float:
        """
        Computes the False Positive Rate (FPR).

        Returns:
            float: FPR score.

        Raises:
            ValueError: If not a binary classification problem.
        """
        if self.num_classes != 2:
            raise ValueError("FPR is only implemented for binary classification.")

        with torch.no_grad():
            cm = self.confusion_matrix()
            fp = cm[0, 1].item()
            tn = cm[0, 0].item()
            fpr = fp / (fp + tn + 1e-12)
            return fpr

    def dice_coefficient(self) -> float:
        """
        Computes the Dice Coefficient.

        Returns:
            float: Dice coefficient score.
        """
        with torch.no_grad():
            if self.num_classes == 2 and self.preds.ndim == 1:
                preds = (torch.sigmoid(self.preds) >= 0.5).long()
                intersection = (preds & self.labels).sum().item()
                total = preds.sum().item() + self.labels.sum().item()
                dice = (2 * intersection) / (total + 1e-12)
                return dice
            else:
                # Multi-class Dice
                cm = self.confusion_matrix()
                intersection = cm.diag()
                total = cm.sum(dim=1) + cm.sum(dim=0)
                dice = (2 * intersection / (total + 1e-12)).mean().item()
                return dice

    def mean_average_precision(self) -> float:
        """
        Computes the Mean Average Precision (mAP).

        Returns:
            float: mAP score.
        """
        # mAP implementation is more suited for object detection.
        # For multi-class classification, it can be interpreted similarly to average precision per class.
        with torch.no_grad():
            average_precisions = []
            probs = self._to_probabilities()
            for cls in range(self.num_classes):
                if self.preds.ndim > 1:
                    prob = probs[:, cls]
                else:
                    prob = probs

                label = (self.labels == cls).long()
                if torch.sum(label) == 0:
                    continue  # Skip classes with no samples
                sorted_prob, sorted_indices = torch.sort(prob, descending=True)
                sorted_label = label[sorted_indices]
                tp = torch.cumsum(sorted_label, dim=0)
                fp = torch.cumsum(1 - sorted_label, dim=0)
                precision = tp / (tp + fp + 1e-12)
                recall = tp / (torch.sum(sorted_label) + 1e-12)
                # Compute AP as area under the precision-recall curve
                average_precision = torch.trapz(precision, recall).item()
                average_precisions.append(average_precision)
            return sum(average_precisions) / len(average_precisions) if average_precisions else 0.0

    def get_all_metrics(self) -> Dict[str, Any]:
        """
        Computes and returns all metrics.

        Returns:
            Dict[str, Any]: A dictionary of metric names and their corresponding scores.
        """
        metrics = {
            'Accuracy': self.accuracy(),
            'Precision': self.precision(),
            'Recall': self.recall(),
            'F1 Score': self.f1_score(),
            'Confusion Matrix': self.confusion_matrix(),
            'ROC-AUC': self.roc_auc() if self.num_classes == 2 else None,
            'PR-AUC': self.pr_auc() if self.num_classes == 2 else None,
            'Top-5 Accuracy': self.top_k_accuracy(k=5) if self.preds.ndim > 1 else None,
            'Log Loss': self.log_loss(),
            'Cross Entropy Loss': self.cross_entropy_loss(),
            'Mean Squared Error': self.mean_squared_error(),
            'Kappa Statistic': self.kappa_score(),
            'Matthews Correlation Coefficient': self.matthews_correlation_coefficient(),
            'Specificity': self.specificity() if self.num_classes == 2 else None,
            'Jaccard Index': self.jaccard_index(),
            'Balanced Accuracy': self.balanced_accuracy(),
            'True Negative Rate': self.true_negative_rate() if self.num_classes == 2 else None,
            'False Positive Rate': self.false_positive_rate() if self.num_classes == 2 else None,
            'Dice Coefficient': self.dice_coefficient(),
            'Mean Average Precision': self.mean_average_precision(),
        }
        # Remove None values for metrics not applicable
        return {k: v for k, v in metrics.items() if v is not None}


# Example usage
if __name__ == "__main__":
    # Sample data for binary classification
    preds_binary = torch.tensor([0.9, 0.2, 0.8, 0.4, 0.6])
    labels_binary = torch.tensor([1, 0, 1, 0, 1])

    metrics_calc_binary = MetricsCalculator(preds=preds_binary, labels=labels_binary, num_classes=2)
    metrics_binary = metrics_calc_binary.get_all_metrics()
    print("Binary Classification Metrics:")
    for metric, value in metrics_binary.items():
        if isinstance(value, Tensor):
            print(f"{metric}: {value.item()}")
        else:
            print(f"{metric}: {value}")

    # Sample data for multi-class classification
    preds_multi = torch.tensor([
        [2.5, 0.3, 2.1],
        [0.1, 3.2, 0.3],
        [1.2, 0.4, 2.2],
        [0.3, 2.8, 0.5],
        [1.1, 0.2, 2.5],
    ])
    labels_multi = torch.tensor([0, 1, 2, 1, 0])

    metrics_calc_multi = MetricsCalculator(preds=preds_multi, labels=labels_multi, num_classes=3)
    metrics_multi = metrics_calc_multi.get_all_metrics()
    print("\nMulti-Class Classification Metrics:")
    for metric, value in metrics_multi.items():
        if isinstance(value, Tensor):
            print(f"{metric}: {value.item()}")
        else:
            print(f"{metric}: {value}")