from __future__ import annotations

from typing import Any, Callable, Dict, Optional

from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    classification_report,
    confusion_matrix,
    balanced_accuracy_score,
    matthews_corrcoef,
    r2_score,
    mean_squared_error,
    mean_absolute_error,
    roc_auc_score,
)

from ..context import Context


class EvaluateModule:
    def __init__(self, context: Context) -> None:
        self._ctx = context

    # ---------------------------------------------------------
    # Internal helpers
    # ---------------------------------------------------------

    def _y(self) -> tuple[Any, Any]:
        if self._ctx.get_predictions() is None:
            raise RuntimeError(
                "No predictions found. Run project.model.predict() before evaluation."
            )

        _, _, _, y_test = self._ctx.get_split_data()
        return y_test, self._ctx.get_predictions()

    def _y_proba(self) -> Optional[Any]:
        model = self._ctx.get_model()
        if hasattr(model, "predict_proba"):
            X_test = self._ctx.get_split_data()[1]
            return model.predict_proba(X_test)[:, 1]
        return None

    # ---------------------------------------------------------
    # Classification
    # ---------------------------------------------------------

    def classification(self, mode: str = "summary") -> Dict[str, Any] | None:
        """
        mode options:
        - "summary" (default): clean metric dict
        - "full": includes report + confusion matrix
        - "pretty": prints sklearn-style output
        """

        y_true, y_pred = self._y()
        y_proba = self._y_proba()

        metrics: Dict[str, Any] = {
            "accuracy": accuracy_score(y_true, y_pred),
            "precision": precision_score(y_true, y_pred),
            "recall": recall_score(y_true, y_pred),
            "f1": f1_score(y_true, y_pred),
            "balanced_accuracy": balanced_accuracy_score(y_true, y_pred),
            "matthews_corrcoef": matthews_corrcoef(y_true, y_pred),
        }

        if y_proba is not None:
            try:
                metrics["roc_auc"] = roc_auc_score(y_true, y_proba)
            except Exception:
                pass

        if mode == "pretty":
            print("\n📊 Classification Report\n")
            print(classification_report(y_true, y_pred))
            print("Confusion Matrix:")
            print(confusion_matrix(y_true, y_pred))
            return None

        if mode == "full":
            metrics["classification_report"] = classification_report(
                y_true, y_pred, output_dict=True
            )
            metrics["confusion_matrix"] = confusion_matrix(
                y_true, y_pred
            ).tolist()

        self._ctx.set_metrics(metrics)
        return metrics

    # ---------------------------------------------------------
    # Regression
    # ---------------------------------------------------------

    def regression(self, mode: str = "summary") -> Dict[str, Any] | None:
        """
        mode options:
        - "summary": basic metrics
        - "pretty": formatted print
        """

        y_true, y_pred = self._y()

        metrics = {
            "r2": r2_score(y_true, y_pred),
            "mse": mean_squared_error(y_true, y_pred),
            "mae": mean_absolute_error(y_true, y_pred),
        }

        if mode == "pretty":
            print("\n📊 Regression Metrics\n")
            for k, v in metrics.items():
                print(f"{k.upper():<10}: {v:.6f}")
            return None

        self._ctx.set_metrics(metrics)
        return metrics

    # ---------------------------------------------------------
    # Custom metric
    # ---------------------------------------------------------

    def custom_metric(
        self, fn: Callable[[Any, Any], float], name: str = "custom_metric"
    ) -> float:
        y_true, y_pred = self._y()
        value = float(fn(y_true, y_pred))
        current = self._ctx.get_metrics() or {}
        current[name] = value
        self._ctx.set_metrics(current)
        return value

    # ---------------------------------------------------------
    # Compatibility aliases
    # ---------------------------------------------------------

    def evaluate_classifier(self, mode: str = "summary") -> Dict[str, Any] | None:
        return self.classification(mode=mode)

    def evaluate_regressor(self, mode: str = "summary") -> Dict[str, Any] | None:
        return self.regression(mode=mode)