import numpy as np


class ClassificationEvalInfo:
    """A class for evaluating classification models.

    Attributes
    ----------
    true_labels : np.ndarray
        The true labels of the data.
    predicted_labels : np.ndarray
        The predicted labels of the data.
    n_classes : int
        The number of classes in the data.
    confusion_matrix : np.ndarray
        The confusion matrix of the data.
    tp : np.ndarray
        The true positives of the data.
    fp : np.ndarray
        The false positives of the data.
    fn : np.ndarray
        The false negatives of the data.
    tn : np.ndarray
        The true negatives of the data.
    accuracy : np.ndarray
        The accuracy of the data.
    precision : np.ndarray
        The precision of the data.
    sensitivity : np.ndarray
        The sensitivity of the data.
    specificity : np.ndarray
        The specificity of the data.
    f1 : np.ndarray
        The f1 score of the data.
    """

    def __init__(self, *,
                 true_labels: np.ndarray,
                 predicted_labels: np.ndarray,
                 n_classes: int | None = None,
                 ) -> None:
        """Initialize the ClassificationEvalInfo object.

        Arguments
        ---------
        true_labels : np.ndarray
            The true labels of the data.
        predicted_labels : np.ndarray
            The predicted labels of the data.
        n_classes : int, optional
            The number of classes in the data. If not provided, it will be inferred from the true labels.
        """
        self.true_labels: np.ndarray = true_labels
        self.predicted_labels: np.ndarray = predicted_labels
        self.n_classes: int | None = n_classes

        self.confusion_matrix: np.ndarray = self._compute_confusion_matrix()
        self.tp: np.ndarray = np.diag(self.confusion_matrix)
        self.fp: np.ndarray = np.sum(self.confusion_matrix, axis=0) - self.tp
        self.fn: np.ndarray = np.sum(self.confusion_matrix, axis=1) - self.tp
        self.tn: np.ndarray = np.sum(self.confusion_matrix) - (self.tp + self.fp + self.fn)

        self.accuracy: np.ndarray = (self.tp + self.tn) / (self.tp + self.fp + self.fn + self.tn)
        self.precision: np.ndarray = self.tp / (self.tp + self.fp)
        self.sensitivity: np.ndarray = self.tp / (self.tp + self.fn)
        self.specificity: np.ndarray = self.tn / (self.tn + self.fp)
        self.f1: np.ndarray = 2 * (self.precision * self.sensitivity) / (self.precision + self.sensitivity)

    def _compute_confusion_matrix(self) -> np.ndarray:
        classes: int = len(np.unique(self.true_labels)) if self.n_classes is None else self.n_classes

        matrix: np.ndarray = np.zeros((classes, classes), dtype=int)

        for true, pred in zip(self.true_labels, self.predicted_labels):
            matrix[true, pred] += 1

        return matrix

    def __repr__(self) -> str:
        accuracy: str = f"accuracy:{self.accuracy.mean(): .4f}"
        precision: str = f"precision:{self.precision.mean(): .4f}"
        sensitivity: str = f"sensitivity:{self.sensitivity.mean(): .4f}"
        specificity: str = f"specificity:{self.specificity.mean(): .4f}"
        f1: str = f"f1:{self.f1.mean(): .4f}"
        confusion_matrix: str = f"confusion_matrix:\n{self.confusion_matrix}"

        return "\n".join([accuracy, precision, sensitivity, specificity, f1, confusion_matrix])
