"""GeoRFRegressor: geometry-aware random forest for regression."""
import numpy as np
from joblib import Parallel, delayed
from sklearn.metrics import r2_score

from hvrt import HVRT

from ._base import _GeoRFBase
from ._utils import fit_one_tree_reg

_K_NEIGHBORS = 3  # kNN neighbours for synthetic-y interpolation


class GeoRFRegressor(_GeoRFBase):
    """Random forest regressor that generates tree diversity via HVRT expansion.

    Each tree is trained on a completely fresh synthetic dataset drawn from a
    single HVRT density model.  Synthetic target values are assigned via k-NN
    weighted interpolation (in HVRT z-space) from the real (X, y) pairs.

    Parameters
    ----------
    n_estimators : int, default=100
        Number of trees in the ensemble.
    n_samples_per_tree : int, default=500
        Synthetic samples generated per tree.
    max_depth : int, default=6
        Maximum depth of each decision tree.
    min_samples_leaf : int, default=5
        Minimum number of samples in a leaf node.
    max_features : str | int | float | None, default=None
        Feature subsampling per split (same semantics as sklearn).
    bandwidth : float | str, default='auto'
        KDE bandwidth passed to HVRT.  ``'auto'`` lets HVRT select the
        optimal kernel width per partition.
    n_jobs : int | None, default=None
        Parallel tree fitting workers.
    random_state : int | None, default=42
        Base seed.  Tree *i* uses seed ``random_state + i``.
    """

    # ------------------------------------------------------------------
    # Fit
    # ------------------------------------------------------------------

    def fit(self, X, y, feature_types=None):
        """Fit the GeoRF regressor.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        y : array-like of shape (n_samples,)
            Continuous target values.
        feature_types : list[str] | None
            Per-column type: ``'continuous'`` or ``'categorical'``.

        Returns
        -------
        self
        """
        X = np.asarray(X, dtype=float)
        y = np.asarray(y, dtype=float)

        # Keep real data for quality / diversity scoring
        self.X_real_ = X.copy()
        self.y_real_ = y.copy()

        # ---- fit a single HVRT on X ----
        self.hvrt_model_ = HVRT(bandwidth=self.bandwidth, random_state=self._base_seed)
        self.hvrt_model_.fit(X, feature_types=feature_types)

        # Precompute z-space representation of real X once; workers reuse it
        X_real_z = self.hvrt_model_._to_z(X)

        # ---- fit trees in parallel ----
        seeds = [self._base_seed + i for i in range(self.n_estimators)]

        self.trees_ = Parallel(n_jobs=self.n_jobs)(
            delayed(fit_one_tree_reg)(
                seed, self.hvrt_model_, X_real_z, y,
                self.n_samples_per_tree, _K_NEIGHBORS, self._tree_kwargs(),
            )
            for seed in seeds
        )

        return self

    # ------------------------------------------------------------------
    # Predict
    # ------------------------------------------------------------------

    def predict(self, X):
        """Predict target values (average over all trees).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)

        Returns
        -------
        ndarray of shape (n_samples,)
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)
        preds = np.array([tree.predict(X) for tree in self.trees_])  # (n_trees, n_samples)
        return preds.mean(axis=0)

    # ------------------------------------------------------------------
    # Interpretability
    # ------------------------------------------------------------------

    def _score_tree(self, tree):
        """R² for a single tree on the real seed data."""
        y_pred = tree.predict(self.X_real_)
        return r2_score(self.y_real_, y_pred)

    def diversity_score(self):
        """Average pairwise normalised disagreement between trees.

        Disagreement is defined as mean absolute difference of normalised
        predictions.  Normalised = divided by the global prediction range.

        Returns
        -------
        float >= 0
        """
        self._check_is_fitted()
        preds = np.array(
            [tree.predict(self.X_real_) for tree in self.trees_]
        )  # (n_trees, n_samples)

        n_trees = len(self.trees_)
        if n_trees < 2:
            return 0.0

        range_ = preds.max() - preds.min()
        if range_ < 1e-10:
            return 0.0
        preds_norm = preds / range_

        total = 0.0
        count = 0
        for i in range(n_trees):
            for j in range(i + 1, n_trees):
                total += float(np.mean(np.abs(preds_norm[i] - preds_norm[j])))
                count += 1

        return total / count

    def prediction_std(self, X):
        """Per-sample std of predictions across trees.

        Measures how much the ensemble disagrees on each sample — a
        natural uncertainty estimate requiring no extra computation beyond
        what predict() already does.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)

        Returns
        -------
        ndarray of shape (n_samples,)
            Higher values = more uncertain predictions.
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)
        preds = np.array([tree.predict(X) for tree in self.trees_])  # (n_trees, n_samples)
        return preds.std(axis=0)

    def _summary_sample_stats(self, X):
        std = self.prediction_std(X)
        return {
            "prediction_uncertainty": {
                "mean_std": float(std.mean()),
                "max_std":  float(std.max()),
            }
        }
