"""Shared base class for GeoRFClassifier and GeoRFRegressor."""
import numpy as np
from joblib import Parallel, delayed


class _GeoRFBase:
    """Base class with shared parameters and interpretability methods.

    Subclasses implement fit(), predict(), _score_tree(), and diversity_score().
    """

    def __init__(
        self,
        n_estimators=100,
        n_samples_per_tree=500,
        max_depth=6,
        min_samples_leaf=5,
        max_features=None,
        bandwidth="auto",
        n_jobs=None,
        random_state=42,
    ):
        self.n_estimators = n_estimators
        self.n_samples_per_tree = n_samples_per_tree
        self.max_depth = max_depth
        self.min_samples_leaf = min_samples_leaf
        self.max_features = max_features
        self.bandwidth = bandwidth
        self.n_jobs = n_jobs
        self.random_state = random_state

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

    @property
    def _base_seed(self):
        return self.random_state if self.random_state is not None else 42

    def _tree_kwargs(self):
        return dict(
            max_depth=self.max_depth,
            min_samples_leaf=self.min_samples_leaf,
            max_features=self.max_features,
        )

    def _check_is_fitted(self):
        if not hasattr(self, "trees_"):
            raise RuntimeError("Call fit() before using this method.")

    # ------------------------------------------------------------------
    # Abstract interface (implemented by subclasses)
    # ------------------------------------------------------------------

    def fit(self, X, y, feature_types=None):
        raise NotImplementedError

    def predict(self, X):
        raise NotImplementedError

    def _score_tree(self, tree):
        """Return quality score (AUC or R²) for a single tree on real data."""
        raise NotImplementedError

    def diversity_score(self):
        """Return average pairwise disagreement rate across all trees."""
        raise NotImplementedError

    # ------------------------------------------------------------------
    # Interpretability methods (shared)
    # ------------------------------------------------------------------

    def feature_importances(self, feature_names=None):
        """Return impurity-based feature importances averaged over all trees.

        Parameters
        ----------
        feature_names : list[str], optional
            Column names for the features.  Falls back to 'feature_0', 'feature_1', …

        Returns
        -------
        dict
            ``{feature_name: importance}`` sorted by descending importance.
        """
        self._check_is_fitted()
        importances = np.mean(
            [tree.feature_importances_ for tree in self.trees_], axis=0
        )
        n_features = len(importances)
        if feature_names is None:
            feature_names = [f"feature_{i}" for i in range(n_features)]
        result = dict(zip(feature_names, importances))
        return dict(sorted(result.items(), key=lambda x: x[1], reverse=True))

    def tree_quality_scores(self):
        """Per-tree quality score evaluated on the real seed data.

        For classifiers: AUC (binary) or macro-OVR AUC (multiclass).
        For regressors: R².

        Returns
        -------
        ndarray of shape (n_estimators,)
        """
        self._check_is_fitted()
        return np.array([self._score_tree(tree) for tree in self.trees_])

    def feature_importance_stability(self, feature_names=None):
        """Std of feature importances across trees — measures ranking stability.

        A low std relative to the mean importance indicates the feature's
        rank is consistent regardless of which synthetic dataset a tree
        was trained on.

        Parameters
        ----------
        feature_names : list[str] | None

        Returns
        -------
        dict[str, float]
            Per-feature importance std, sorted descending (most variable first).
        """
        self._check_is_fitted()
        all_imp = np.array([t.feature_importances_ for t in self.trees_])
        stds = all_imp.std(axis=0)
        n = len(stds)
        if feature_names is None:
            feature_names = [f"feature_{i}" for i in range(n)]
        result = dict(zip(feature_names, stds))
        return dict(sorted(result.items(), key=lambda x: x[1], reverse=True))

    def _summary_sample_stats(self, X):
        """Subclass hook: return sample-level stats dict given X."""
        return {}

    def summary(self, X=None, feature_names=None):
        """One-call interpretability report.

        Collates provenance, tree quality, diversity, feature importances
        with stability, and — when X is provided — per-sample uncertainty.

        GeoRF has exactly two stochasticity sources (both seeded):
          1. HVRT KDE sampling — controls which synthetic points each tree trains on.
          2. DecisionTree tie-breaking — controls split selection under equal gain.
        This makes every prediction fully reproducible and auditable.

        Parameters
        ----------
        X : array-like | None
            If provided, per-sample uncertainty stats are included.
        feature_names : list[str] | None

        Returns
        -------
        dict
        """
        self._check_is_fitted()
        quality = self.tree_quality_scores()
        result = {
            "model_type": type(self).__name__,
            "n_estimators": self.n_estimators,
            "n_features": self.X_real_.shape[1],
            "stochasticity": {
                "n_sources": 2,
                "description": (
                    "HVRT KDE sampling (bandwidth='auto') + "
                    "DecisionTree tie-breaking — both seeded, fully reproducible"
                ),
                "base_seed": self._base_seed,
            },
            "provenance": self.provenance(),
            "tree_quality": {
                "mean": float(quality.mean()),
                "std":  float(quality.std()),
                "min":  float(quality.min()),
                "max":  float(quality.max()),
            },
            "diversity_score": self.diversity_score(),
            "feature_importances": self.feature_importances(feature_names),
            "feature_importance_stability": self.feature_importance_stability(feature_names),
        }
        if X is not None:
            result.update(self._summary_sample_stats(np.asarray(X, dtype=float)))
        return result

    def provenance(self):
        """Return metadata about the fit.

        Returns
        -------
        dict with keys:
            n_real_samples, n_synthetic_per_tree, n_total_synthetic,
            n_estimators, expansion_ratio
        """
        self._check_is_fitted()
        n_real = len(self.X_real_)
        return {
            "n_real_samples": n_real,
            "n_synthetic_per_tree": self.n_samples_per_tree,
            "n_total_synthetic": self.n_estimators * self.n_samples_per_tree,
            "n_estimators": self.n_estimators,
            "expansion_ratio": self.n_samples_per_tree / n_real,
        }
