import numpy as np

from ..log import ansi


def _safe_div(numerator: np.ndarray | float, denominator: np.ndarray | float) -> np.ndarray:
    num = np.asarray(numerator, dtype=float)
    den = np.asarray(denominator, dtype=float)
    return np.divide(num, den, out=np.zeros_like(num, dtype=float), where=den != 0)


class EvalInfo:
    """A class that calculates and stores evaluation metrics for regression and classification tasks.

    Attributes
    ----------
    self.task : str
        The task type, either 'regression' or 'classification'.
    self.true_output : np.ndarray
        The true output values.
    self.predicted_output : np.ndarray
        The predicted output values.
    self.predicted_std : np.ndarray
        The standard deviation of the predicted output values.
    self.mae : float
        The mean absolute error. Only available for regression tasks.
    self.mse : float
        The mean squared error. Only available for regression tasks.
    self.rmse : float
        The root mean squared error. Only available for regression tasks.
    self.r2 : float
        The R2 score. Only available for regression tasks.
    self.confusion_matrix : np.ndarray
        The confusion matrix. Only available for classification tasks.
    self.tp : np.ndarray
        The true positive values. Only available for classification tasks.
    self.fp : np.ndarray
        The false positive values. Only available for classification tasks.
    self.fn : np.ndarray
        The false negative values. Only available for classification tasks.
    self.tn : np.ndarray
        The true negative values. Only available for classification tasks.
    self.accuracy : float
        The accuracy score. Only available for classification tasks.
    self.precision : float
        The precision score. Only available for classification tasks.
    self.recall : float
        The recall score. Only available for classification tasks.
    self.sensitivity : float
        The sensitivity score. Only available for classification tasks.
    self.specificity : float
        The specificity score. Only available for classification tasks.
    self.f1_score : float
        The F1 score. Only available for classification tasks.
    """

    def __init__(self,
                 task: str,
                 true_output: np.ndarray,
                 predicted_output: np.ndarray,
                 predicted_std: np.ndarray | None = None,
                 n_classes: int | None = None
                 ) -> None:
        """
        Parameters
        ----------
        task : str
            The task type, either 'regression' or 'classification'.
        true_output : np.ndarray
            The true output values.
        predicted_output : np.ndarray
            The predicted output values.
        predicted_std : np.ndarray, optional
            The standard deviation of the predicted output values. If not provided, it is set to an array of zeros.
        n_classes : int, optional
            The number of classes in the classification task. If not provided, it is set to the number of unique classes.
        """

        self.task = task
        self.true_output = true_output
        self.predicted_output = predicted_output
        self.predicted_std = predicted_std if predicted_std is not None else np.zeros_like(predicted_output)

        if self.task == 'regression':
            self._calculate_regression_metrics(self.true_output, self.predicted_output)
        elif self.task in ('classification', 'signal_classification'):
            self._calculate_classification_metrics(self.true_output, self.predicted_output, n_classes)
        else:
            raise ValueError('Invalid task type')

    def _calculate_regression_metrics(self, y: np.ndarray, y_pred: np.ndarray) -> None:
        y = np.asarray(y, dtype=float)
        y_pred = np.asarray(y_pred, dtype=float)

        self.mae = float(np.mean(np.abs(y - y_pred)))
        self.mse = float(np.mean((y - y_pred) ** 2))
        self.rmse = float(np.sqrt(self.mse))

        sse = float(np.sum((y - y_pred) ** 2))
        sst = float(np.sum((y - np.mean(y)) ** 2))
        self.r2 = float(1.0 - sse / sst) if sst != 0.0 else 0.0

    def _calculate_classification_metrics(self, y: np.ndarray, y_pred: np.ndarray, n_classes: int | None) -> None:
        y_true = np.asarray(y).reshape(-1).astype(int)
        y_pred = np.asarray(y_pred)
        y_pred_labels = np.argmax(y_pred, axis=1).astype(int) if y_pred.ndim > 1 else y_pred.astype(int).reshape(-1)

        self.confusion_matrix = self._compute_confusion_matrix(y_true, y_pred_labels, n_classes)

        self.tp = np.diag(self.confusion_matrix).astype(float)
        self.fp = np.sum(self.confusion_matrix, axis=0).astype(float) - self.tp
        self.fn = np.sum(self.confusion_matrix, axis=1).astype(float) - self.tp
        self.tn = float(np.sum(self.confusion_matrix)) - (self.tp + self.fp + self.fn)

        precision_per_class = _safe_div(self.tp, self.tp + self.fp)
        recall_per_class = _safe_div(self.tp, self.tp + self.fn)
        specificity_per_class = _safe_div(self.tn, self.tn + self.fp)
        f1_per_class = _safe_div(2 * precision_per_class * recall_per_class, precision_per_class + recall_per_class)

        self.accuracy = float(np.mean(y_true == y_pred_labels)) if y_true.size else 0.0
        self.precision = float(np.mean(precision_per_class))
        self.recall = self.sensitivity = float(np.mean(recall_per_class))
        self.specificity = float(np.mean(specificity_per_class))
        self.f1_score = float(np.mean(f1_per_class))

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

        matrix: np.ndarray = np.zeros((classes, classes), dtype=int)
        for true, pred in zip(y, y_pred):
            if true < 0 or true >= classes or pred < 0 or pred >= classes:
                raise ValueError(f"found label out of range [0, {classes - 1}]")
            matrix[true, pred] += 1

        return matrix

    def _regression_repr(self) -> str:
        return f"\n{ansi.BOLD}{ansi.BLUE}-> regression metrics{ansi.RESET}\n" + \
            f"   |> MAE:  {self.mae:.4f}\n" + \
            f"   |> MSE:  {self.mse:.4f}\n" + \
            f"   |> RMSE: {self.rmse:.4f}\n" + \
            f"   |> R2:   {self.r2:.4f}"

    def _classification_repr(self) -> str:
        return f"\n{ansi.BOLD}{ansi.BLUE}-> classification metrics{ansi.RESET}\n" + \
            f"   |> accuracy:    {self.accuracy:.4f}\n" + \
            f"   |> precision:   {self.precision:.4f}\n" + \
            f"   |> recall:      {self.recall:.4f}\n" + \
            f"   |> sensitivity: {self.sensitivity:.4f}\n" + \
            f"   |> specificity: {self.specificity:.4f}\n" + \
            f"   |> f1 score:    {self.f1_score:.4f}"

    def __repr__(self) -> str:
        if self.task == 'regression':
            return self._regression_repr()
        elif self.task in ('classification', 'signal_classification'):
            return self._classification_repr()
        else:
            raise ValueError('Invalid task type')
