import numpy as np
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    matthews_corrcoef,
    cohen_kappa_score,
    log_loss,
    roc_auc_score,
    confusion_matrix,
)

from .base import BaseEvaluator


class ClassificationEvaluator(BaseEvaluator):
    """
    Evaluator for classification models.
    Supports binary and multi-class classification.
    """

    def evaluate(self) -> dict:
        if self.model_type != "classification":
            raise ValueError("Provided model is not a classification model")

        y_pred = self.model.predict(self.X_test)

        y_true = np.asarray(self.y_test)
        y_pred = np.asarray(y_pred)

        labels = np.unique(np.concatenate([y_true, y_pred]))
        is_binary = len(labels) == 2

        # ── Core metrics ────────────────────────────────
        metrics = {
            "accuracy": float(accuracy_score(y_true, y_pred)),
            "balanced_accuracy": float(balanced_accuracy_score(y_true, y_pred)),
            "precision": float(
                precision_score(y_true, y_pred, average="weighted", zero_division=0)
            ),
            "recall": float(
                recall_score(y_true, y_pred, average="weighted", zero_division=0)
            ),
            "f1_score": float(
                f1_score(y_true, y_pred, average="weighted", zero_division=0)
            ),
            "mcc": float(matthews_corrcoef(y_true, y_pred)),
            "cohen_kappa": float(cohen_kappa_score(y_true, y_pred)),
        }

        # ── Probability-based metrics (if available) ────
        has_proba = callable(getattr(self.model, "predict_proba", None))

        if has_proba:
            y_proba = self.model.predict_proba(self.X_test)
            metrics["log_loss"] = float(log_loss(y_true, y_proba, labels=labels))

            if is_binary:
                metrics["roc_auc"] = float(
                    roc_auc_score(y_true, y_proba[:, 1])
                )
            else:
                try:
                    metrics["roc_auc"] = float(
                        roc_auc_score(
                            y_true, y_proba,
                            multi_class="ovr", average="weighted",
                        )
                    )
                except ValueError:
                    pass  # not enough classes present

        # ── Confusion matrix ────────────────────────────
        cm = confusion_matrix(y_true, y_pred, labels=labels)

        self.confusion_matrix_ = cm
        self.class_labels_ = labels.tolist()

        return metrics