"""Ensemble pruning for GeoRFRegressor.

Implements greedy backward elimination: trees are removed one at a time if
doing so improves R² on a held-out validation set.  The algorithm terminates
when no single removal yields an improvement, or when the ensemble would fall
below ``min_trees``.

Complexity: O(N²) evaluations in the worst case, where N is the initial
number of trees.  Each evaluation is a mean of (N - k) tree predictions over
the validation set.  At N=100 trees and ~90 validation samples this completes
in well under a second.
"""
import numpy as np
from sklearn.metrics import r2_score


def greedy_prune(model, X_val, y_val, min_trees=5, verbose=True):
    """Greedy backward elimination of trees from a GeoRFRegressor.

    Iteratively removes the single tree whose removal maximally improves R²
    on the provided validation data.  Stops when no removal helps or fewer
    than ``min_trees`` trees would remain.

    Modifies ``model.trees_`` **in place** and attaches ``model.pruning_history_``.

    Parameters
    ----------
    model : GeoRFRegressor
        A fitted regressor with a ``trees_`` attribute.
    X_val : array-like of shape (n_val, n_features)
        Held-out validation features used to guide pruning.  Must **not** be
        the same data used to fit the model (would overfit).
    y_val : array-like of shape (n_val,)
        Held-out validation targets.
    min_trees : int, default=5
        Minimum number of trees to retain.
    verbose : bool, default=True
        Print a line for each removal step.

    Returns
    -------
    history : list of (n_trees_remaining, val_r2) tuples
        One entry per pruning step, including the initial state.

    Raises
    ------
    ValueError
        If ``model.trees_`` is empty or ``min_trees`` is not positive.
    """
    if not hasattr(model, "trees_") or len(model.trees_) == 0:
        raise ValueError("model.trees_ is empty or missing.  Fit the model first.")
    if min_trees < 1:
        raise ValueError("min_trees must be >= 1.")

    X_val = np.asarray(X_val, dtype=float)
    y_val = np.asarray(y_val, dtype=float)

    trees = list(model.trees_)

    def _score(subset):
        preds = np.mean([t.predict(X_val) for t in subset], axis=0)
        return r2_score(y_val, preds)

    current_score = _score(trees)
    history = [(len(trees), current_score)]

    if verbose:
        print(f"  Initial: {len(trees):3d} trees, val R2={current_score:.4f}")

    improved = True
    while improved and len(trees) > min_trees:
        improved = False
        best_score = current_score
        best_idx = -1

        for i in range(len(trees)):
            candidate = trees[:i] + trees[i + 1:]
            s = _score(candidate)
            if s > best_score:
                best_score = s
                best_idx = i

        if best_idx >= 0:
            removed = trees.pop(best_idx)
            current_score = best_score
            history.append((len(trees), current_score))
            improved = True
            if verbose:
                print(
                    f"  Removed tree {best_idx:3d}: "
                    f"{len(trees):3d} trees remaining, val R2={current_score:.4f}"
                )

    if verbose:
        n_removed = history[0][0] - history[-1][0]
        delta = history[-1][1] - history[0][1]
        print(
            f"  Done: removed {n_removed} trees, "
            f"val R2 {history[0][1]:.4f} -> {history[-1][1]:.4f} "
            f"({delta:+.4f})"
        )

    model.trees_ = trees
    model.pruning_history_ = history
    return history
