"""Feature importance: impurity-based (MDI), surrogate-based, and permutation."""

import numpy as np
from sklearn.metrics import accuracy_score, r2_score


def impurity_importance(estimator):
    """Compute impurity-based (MDI) feature importance.

    Sum of weighted gain at each node, by feature. For ensembles,
    averages across trees.
    """
    trees = _get_trees(estimator)
    n_features = estimator.n_features_in_

    total_importance = np.zeros(n_features, dtype=np.float64)

    for tree_root in trees:
        tree_imp = np.zeros(n_features, dtype=np.float64)
        _accumulate_impurity(tree_root, tree_imp)
        total_importance += tree_imp

    total_importance /= len(trees)

    # Normalize
    s = total_importance.sum()
    if s > 0:
        total_importance /= s

    return total_importance


def _accumulate_impurity(node, importance):
    """Recursively accumulate weighted gain per feature."""
    if node.is_leaf or node.split is None:
        return

    importance[node.split.feature] += node.weighted_n_samples * node.split.gain

    _accumulate_impurity(node.left_child, importance)
    _accumulate_impurity(node.right_child, importance)


def surrogate_importance(estimator):
    """Compute surrogate-based feature importance.

    For each surrogate appearance in the tree, accumulates
    weighted_n_samples * lambda. This captures feature redundancy
    and substitutability — a feature scores high if it frequently
    serves as a good surrogate for primary splits.
    """
    trees = _get_trees(estimator)
    n_features = estimator.n_features_in_

    total_importance = np.zeros(n_features, dtype=np.float64)

    for tree_root in trees:
        tree_imp = np.zeros(n_features, dtype=np.float64)
        _accumulate_surrogate(tree_root, tree_imp)
        total_importance += tree_imp

    total_importance /= len(trees)

    s = total_importance.sum()
    if s > 0:
        total_importance /= s

    return total_importance


def _accumulate_surrogate(node, importance):
    """Recursively accumulate surrogate-based importance."""
    if node.is_leaf or node.split is None:
        return

    # Primary feature gets its impurity contribution
    importance[node.split.feature] += node.weighted_n_samples * node.split.gain

    # Surrogates get weighted lambda contribution
    for surr in node.split.surrogates:
        importance[surr.feature] += node.weighted_n_samples * surr.lambda_

    _accumulate_surrogate(node.left_child, importance)
    _accumulate_surrogate(node.right_child, importance)


def permutation_importance(estimator, X, y, n_repeats=10, random_state=None, scoring=None):
    """Compute permutation importance, correctly handling NaN patterns.

    When permuting a feature, NaN positions are permuted along with values,
    preserving the overall missingness rate but breaking the feature-target
    association.

    Parameters
    ----------
    estimator : fitted estimator
    X : array-like of shape (n_samples, n_features)
    y : array-like of shape (n_samples,)
    n_repeats : int
    random_state : int or None
    scoring : callable or None (uses accuracy for clf, R2 for reg)

    Returns
    -------
    dict with 'importances_mean', 'importances_std', 'importances'
    """
    rng = np.random.RandomState(random_state)
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y)
    n_features = X.shape[1]

    if scoring is None:
        is_clf = hasattr(estimator, "classes_")
        if is_clf:
            scoring = lambda est, X, y: accuracy_score(y, est.predict(X))
        else:
            scoring = lambda est, X, y: r2_score(y, est.predict(X))

    baseline_score = scoring(estimator, X, y)

    importances = np.zeros((n_features, n_repeats), dtype=np.float64)

    for f in range(n_features):
        for r in range(n_repeats):
            X_perm = X.copy()
            perm_idx = rng.permutation(X.shape[0])
            X_perm[:, f] = X[perm_idx, f]  # NaN positions get shuffled too
            perm_score = scoring(estimator, X_perm, y)
            importances[f, r] = baseline_score - perm_score

    return {
        "importances_mean": importances.mean(axis=1),
        "importances_std": importances.std(axis=1),
        "importances": importances,
    }


def _get_trees(estimator):
    """Extract tree roots from an estimator (single tree or forest)."""
    if hasattr(estimator, "estimators_"):
        return estimator.estimators_
    elif hasattr(estimator, "tree_"):
        return [estimator.tree_]
    else:
        raise ValueError("Estimator must have `tree_` or `estimators_` attribute")
