"""Feature contribution decomposition for GeoRF models.

Provides exact per-sample, per-feature contribution breakdowns:

- ``tree_contributions``: Saabas exact path decomposition for GeoRFRegressor.
  For every sample the prediction decomposes as:
      prediction = bias + sum_f(contribution_f)
  where bias is the ensemble-averaged root-node value and each contribution_f
  is the total change in node value attributed to feature f along the
  root-to-leaf path, averaged across all trees.

- ``xgb_contributions``: SHAP Tree values for GeoXGBRegressor, obtained via
  XGBoost's native Booster.predict(pred_contribs=True).  Same output format
  as ``tree_contributions`` for easy comparison.
"""
import numpy as np


# ---------------------------------------------------------------------------
# Internal: single-tree Saabas decomposition
# ---------------------------------------------------------------------------

def _single_tree_contributions(tree, X):
    """Compute Saabas contributions for one DecisionTreeRegressor.

    Parameters
    ----------
    tree : sklearn.tree.DecisionTreeRegressor
        A fitted tree.
    X : ndarray of shape (n_samples, n_features)

    Returns
    -------
    bias : float
        Root-node mean value (the tree's unconditional prediction).
    contributions : ndarray of shape (n_samples, n_features)
        Per-sample, per-feature cumulative delta along the decision path.
        ``bias + contributions.sum(axis=1)`` equals ``tree.predict(X)``
        exactly (within floating-point tolerance).
    """
    tree_ = tree.tree_
    n_samples, n_features = X.shape
    node_values = tree_.value[:, 0, 0]  # regression: single output, single class

    # decision_path: CSR sparse matrix (n_samples, n_nodes)
    # entry [i, j] is non-zero if sample i passes through node j
    indicator = tree.decision_path(X).toarray().astype(bool)  # dense bool

    contributions = np.zeros((n_samples, n_features))

    for node_idx in range(tree_.node_count):
        feat = tree_.feature[node_idx]
        if feat < 0:  # leaf node (TREE_LEAF = -1, TREE_UNDEFINED = -2)
            continue

        left  = tree_.children_left[node_idx]
        right = tree_.children_right[node_idx]

        at_node    = indicator[:, node_idx]
        went_left  = at_node & indicator[:, left]
        went_right = at_node & indicator[:, right]

        contributions[went_left,  feat] += node_values[left]  - node_values[node_idx]
        contributions[went_right, feat] += node_values[right] - node_values[node_idx]

    return float(node_values[0]), contributions


# ---------------------------------------------------------------------------
# Public: GeoRFRegressor
# ---------------------------------------------------------------------------

def tree_contributions(model, X, feature_names=None):
    """Exact Saabas path contributions for a GeoRFRegressor.

    Decomposes each prediction as ``bias + sum_f(contribution_f)`` where
    *bias* is the grand mean across all trees and *contribution_f* is the
    total signed change in node value attributed to feature *f* along the
    root-to-leaf path, averaged over the full ensemble.

    Parameters
    ----------
    model : GeoRFRegressor
        A fitted regressor.  Must have a ``trees_`` attribute.
    X : array-like of shape (n_samples, n_features)
    feature_names : list[str] | None
        If None, names default to ``'feature_0'``, ``'feature_1'``, ...

    Returns
    -------
    dict with keys:
        ``'bias'`` : float
            Ensemble-averaged root-node value.
        ``'contributions'`` : ndarray of shape (n_samples, n_features)
            Signed feature contributions.  Positive = pushed prediction up.
        ``'feature_names'`` : list[str]
        ``'predictions'`` : ndarray of shape (n_samples,)
            ``bias + contributions.sum(axis=1)`` — equals ``model.predict(X)``
            within numerical tolerance.

    Raises
    ------
    RuntimeError
        If the decomposition is not self-consistent.
    """
    if not hasattr(model, "trees_"):
        raise RuntimeError(
            "tree_contributions requires a fitted GeoRFRegressor "
            "(model.trees_ not found)."
        )
    X = np.asarray(X, dtype=float)
    n_samples, n_features = X.shape

    if feature_names is None:
        feature_names = [f"feature_{i}" for i in range(n_features)]

    biases = []
    all_contributions = []

    for tree in model.trees_:
        bias, contribs = _single_tree_contributions(tree, X)
        biases.append(bias)
        all_contributions.append(contribs)

    mean_bias = float(np.mean(biases))
    mean_contributions = np.mean(all_contributions, axis=0)  # (n_samples, n_features)
    predictions = mean_bias + mean_contributions.sum(axis=1)

    # Consistency check against model.predict (should agree within float precision)
    model_preds = model.predict(X)
    max_err = np.abs(predictions - model_preds).max()
    if max_err > 1e-4:
        raise RuntimeError(
            f"Contribution decomposition is inconsistent with model.predict "
            f"(max error {max_err:.2e}).  This is a bug."
        )

    return {
        "bias": mean_bias,
        "contributions": mean_contributions,
        "feature_names": list(feature_names),
        "predictions": predictions,
    }


# ---------------------------------------------------------------------------
# Internal: single-tree Saabas decomposition for classifiers
# ---------------------------------------------------------------------------

def _single_tree_contributions_clf(tree, X, class_idx):
    """Saabas contributions for one DecisionTreeClassifier.

    Decomposes the predicted probability of ``class_idx`` as
    ``bias + sum_f(contribution_f)``.

    Returns
    -------
    bias : float
        Root-node probability for ``class_idx``.
    contributions : ndarray of shape (n_samples, n_features)
    """
    tree_ = tree.tree_
    n_samples, n_features = X.shape

    # Node class probabilities: value has shape (n_nodes, 1, n_classes)
    counts = tree_.value[:, 0, :].astype(float)          # (n_nodes, n_classes)
    totals = counts.sum(axis=1, keepdims=True)
    proba  = counts / np.where(totals > 0, totals, 1.0)  # row-normalised
    node_p = proba[:, class_idx]                          # (n_nodes,)

    indicator = tree.decision_path(X).toarray().astype(bool)
    contributions = np.zeros((n_samples, n_features))

    for node_idx in range(tree_.node_count):
        feat = tree_.feature[node_idx]
        if feat < 0:
            continue
        left  = tree_.children_left[node_idx]
        right = tree_.children_right[node_idx]

        at_node    = indicator[:, node_idx]
        went_left  = at_node & indicator[:, left]
        went_right = at_node & indicator[:, right]

        contributions[went_left,  feat] += node_p[left]  - node_p[node_idx]
        contributions[went_right, feat] += node_p[right] - node_p[node_idx]

    return float(node_p[0]), contributions


# ---------------------------------------------------------------------------
# Public: GeoRFClassifier
# ---------------------------------------------------------------------------

def clf_contributions(model, X, feature_names=None, class_idx=1):
    """Exact Saabas path contributions for a GeoRFClassifier.

    Decomposes the ensemble's predicted probability for ``class_idx`` as
    ``bias + sum_f(contribution_f)``.  The decomposition is exact —
    ``bias + contributions.sum(axis=1)`` equals
    ``model.predict_proba(X)[:, class_idx]`` within floating-point tolerance.

    Parameters
    ----------
    model : GeoRFClassifier
        A fitted classifier.
    X : array-like of shape (n_samples, n_features)
    feature_names : list[str] | None
    class_idx : int, default=1
        Which class probability to decompose.  For binary classification
        ``class_idx=1`` decomposes the positive-class probability.

    Returns
    -------
    dict with keys:
        ``'bias'`` : float
        ``'contributions'`` : ndarray of shape (n_samples, n_features)
        ``'feature_names'`` : list[str]
        ``'predictions'`` : ndarray of shape (n_samples,)
            Equals ``model.predict_proba(X)[:, class_idx]`` within tolerance.
        ``'class_idx'`` : int
        ``'class_name'`` : str

    Raises
    ------
    RuntimeError
        If the decomposition is not self-consistent.
    """
    if not hasattr(model, "trees_"):
        raise RuntimeError(
            "clf_contributions requires a fitted GeoRFClassifier "
            "(model.trees_ not found)."
        )
    X = np.asarray(X, dtype=float)
    n_samples, n_features = X.shape

    if feature_names is None:
        feature_names = [f"feature_{i}" for i in range(n_features)]

    biases, all_contributions = [], []
    for tree in model.trees_:
        bias, contribs = _single_tree_contributions_clf(tree, X, class_idx)
        biases.append(bias)
        all_contributions.append(contribs)

    mean_bias          = float(np.mean(biases))
    mean_contributions = np.mean(all_contributions, axis=0)
    predictions        = mean_bias + mean_contributions.sum(axis=1)

    model_proba = model.predict_proba(X)[:, class_idx]
    max_err = np.abs(predictions - model_proba).max()
    if max_err > 1e-4:
        raise RuntimeError(
            f"Contribution decomposition inconsistent with predict_proba "
            f"(max error {max_err:.2e}).  This is a bug."
        )

    return {
        "bias":          mean_bias,
        "contributions": mean_contributions,
        "feature_names": list(feature_names),
        "predictions":   predictions,
        "class_idx":     class_idx,
        "class_name":    str(model.classes_[class_idx]),
    }


# ---------------------------------------------------------------------------
# Public: GeoXGBRegressor
# ---------------------------------------------------------------------------

def xgb_contributions(model, X, feature_names=None):
    """SHAP Tree contributions for a GeoXGBRegressor.

    Uses XGBoost's native ``Booster.predict(pred_contribs=True)`` which
    returns exact SHAP values for tree ensembles.  Output format matches
    ``tree_contributions`` for easy comparison.

    Parameters
    ----------
    model : GeoXGBRegressor
        A fitted model.  Must have a ``model_`` attribute (XGBRegressor).
    X : array-like of shape (n_samples, n_features)
    feature_names : list[str] | None

    Returns
    -------
    dict with same keys as ``tree_contributions``:
        ``'bias'``, ``'contributions'``, ``'feature_names'``, ``'predictions'``
    """
    if not hasattr(model, "model_"):
        raise RuntimeError(
            "xgb_contributions requires a fitted GeoXGBRegressor "
            "(model.model_ not found)."
        )
    try:
        import xgboost as xgb
    except ImportError as e:
        raise ImportError(
            "xgboost is required for xgb_contributions.  "
            "Install with: pip install xgboost"
        ) from e

    X = np.asarray(X, dtype=float)
    n_samples, n_features = X.shape

    if feature_names is None:
        feature_names = [f"feature_{i}" for i in range(n_features)]

    booster = model.model_.get_booster()
    dmatrix = xgb.DMatrix(X)
    # pred_contribs=True returns (n_samples, n_features + 1)
    # last column is the expected value (bias)
    raw = booster.predict(dmatrix, pred_contribs=True)

    contributions = raw[:, :n_features].astype(float)
    bias = float(raw[:, n_features].mean())  # constant across rows (same base margin)
    predictions = bias + contributions.sum(axis=1)

    return {
        "bias": bias,
        "contributions": contributions,
        "feature_names": list(feature_names),
        "predictions": predictions,
    }
