"""Pure metric computation functions — all return JSON-serializable dicts."""
import numpy as np
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
    confusion_matrix,
    r2_score,
    mean_squared_error,
    mean_absolute_error,
)


def compute_classification_metrics(y_true, y_pred, y_proba=None, classes=None):
    """Compute classification metrics.

    Parameters
    ----------
    y_true : array-like
    y_pred : array-like
    y_proba : array-like of shape (n_samples, n_classes), optional
    classes : array-like, optional
        Class labels. Inferred from y_true if None.

    Returns
    -------
    dict with keys:
        accuracy, f1_weighted, auc_macro (if y_proba given),
        per_class (precision/recall/f1 per class), confusion_matrix
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    if classes is None:
        classes = np.unique(np.concatenate([y_true, y_pred]))

    result = {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "f1_weighted": float(f1_score(y_true, y_pred, average="weighted", zero_division=0)),
    }

    if y_proba is not None:
        y_proba = np.asarray(y_proba)
        try:
            if len(classes) == 2:
                auc = roc_auc_score(y_true, y_proba[:, 1])
            else:
                auc = roc_auc_score(y_true, y_proba, multi_class="ovr", average="macro")
            result["auc_macro"] = float(auc)
        except Exception:
            result["auc_macro"] = None

    precision = precision_score(y_true, y_pred, average=None, labels=classes, zero_division=0)
    recall = recall_score(y_true, y_pred, average=None, labels=classes, zero_division=0)
    f1 = f1_score(y_true, y_pred, average=None, labels=classes, zero_division=0)

    result["per_class"] = {
        str(cls): {
            "precision": float(precision[i]),
            "recall": float(recall[i]),
            "f1": float(f1[i]),
        }
        for i, cls in enumerate(classes)
    }

    cm = confusion_matrix(y_true, y_pred, labels=classes)
    result["confusion_matrix"] = cm.tolist()

    return result


def compute_regression_metrics(y_true, y_pred):
    """Compute regression metrics.

    Parameters
    ----------
    y_true : array-like
    y_pred : array-like

    Returns
    -------
    dict with keys:
        r2, rmse, mae, mape (masked for y_true==0),
        residuals (mean, std, min, max, p25, p75)
    """
    y_true = np.asarray(y_true, dtype=float)
    y_pred = np.asarray(y_pred, dtype=float)
    residuals = y_pred - y_true

    result = {
        "r2": float(r2_score(y_true, y_pred)),
        "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
        "mae": float(mean_absolute_error(y_true, y_pred)),
    }

    # MAPE — mask zeros to avoid division by zero
    nonzero = y_true != 0
    if nonzero.any():
        result["mape"] = float(np.mean(np.abs((y_true[nonzero] - y_pred[nonzero]) / y_true[nonzero])))
    else:
        result["mape"] = None

    result["residuals"] = {
        "mean": float(residuals.mean()),
        "std": float(residuals.std()),
        "min": float(residuals.min()),
        "max": float(residuals.max()),
        "p25": float(np.percentile(residuals, 25)),
        "p75": float(np.percentile(residuals, 75)),
    }

    return result


def compute_georf_metrics(estimator, feature_names=None):
    """Compute GeoRF-specific interpretability metrics via duck-typing.

    Only calls methods/attributes that exist on the estimator; works on any
    sklearn-compatible model but returns richer results for GeoRF instances.

    Parameters
    ----------
    estimator : fitted estimator
    feature_names : list[str], optional

    Returns
    -------
    dict (may be empty for non-GeoRF estimators)
    """
    result = {}

    if hasattr(estimator, "diversity_score"):
        try:
            result["diversity_score"] = float(estimator.diversity_score())
        except Exception:
            pass

    if hasattr(estimator, "tree_quality_scores"):
        try:
            scores = np.asarray(estimator.tree_quality_scores())
            result["tree_quality_mean"] = float(scores.mean())
            result["tree_quality_std"] = float(scores.std())
            result["tree_quality_min"] = float(scores.min())
            result["tree_quality_max"] = float(scores.max())
        except Exception:
            pass

    if hasattr(estimator, "provenance"):
        try:
            result["provenance"] = estimator.provenance()
        except Exception:
            pass

    if hasattr(estimator, "feature_importances"):
        try:
            fi = estimator.feature_importances(feature_names=feature_names)
            result["feature_importances"] = {k: float(v) for k, v in fi.items()}
        except Exception:
            pass
    elif hasattr(estimator, "feature_importances_"):
        try:
            importances = np.asarray(estimator.feature_importances_)
            n = len(importances)
            names = feature_names if feature_names is not None else [f"feature_{i}" for i in range(n)]
            result["feature_importances"] = {
                k: float(v)
                for k, v in sorted(zip(names, importances), key=lambda x: x[1], reverse=True)
            }
        except Exception:
            pass

    return result
