import numpy as np
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
)

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)

        return {
            "accuracy": float(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)
            ),
        }