"""GeoXGBRegressor: Geometry-Aware XGBoost Regressor with HVRT augmentation.

Unlike GeoRFRegressor, which trains N decision trees each on a *different*
synthetic dataset, GeoXGBRegressor performs a single HVRT expansion and trains
one XGBoost ensemble on the concatenated real + synthetic data.  XGBoost is
already an ensemble internally (gradient-boosted trees), so per-tree synthetic
diversity is not needed — the benefit comes from enriching the training
distribution before boosting begins.
"""
import numpy as np

try:
    from xgboost import XGBRegressor as _XGBRegressor
except ImportError as _err:
    raise ImportError(
        "xgboost is required for GeoXGBRegressor. "
        "Install it with: pip install xgboost"
    ) from _err

from hvrt import HVRT

from ._utils import _make_synthetic_dataset


class GeoXGBRegressor:
    """XGBoost regressor trained on real + HVRT-augmented synthetic data.

    Parameters
    ----------
    n_synthetic : int, default=1000
        Number of synthetic samples generated by HVRT at fit time.
        Combined with all real samples as training data for XGBoost.
    bandwidth : float | str, default='auto'
        KDE bandwidth passed to HVRT.  ``'auto'`` lets HVRT select the
        optimal kernel width per partition.
    n_estimators : int, default=300
        XGBoost boosting rounds.
    max_depth : int, default=6
        XGBoost maximum tree depth.
    learning_rate : float, default=0.05
        XGBoost step-size shrinkage (eta).
    subsample : float, default=0.8
        Row subsampling ratio per boosting round.
    colsample_bytree : float, default=0.8
        Column subsampling ratio per tree.
    min_child_weight : int, default=3
        Minimum sum of instance weight in a child node.
    reg_alpha : float, default=0.0
        L1 regularisation on leaf weights.
    reg_lambda : float, default=1.0
        L2 regularisation on leaf weights.
    n_jobs : int | None, default=None
        Passed to XGBRegressor.  None uses XGBoost's default threading.
    random_state : int | None, default=42
        Seed for both HVRT expansion and XGBoost.
    """

    def __init__(
        self,
        n_synthetic=1000,
        bandwidth="auto",
        n_estimators=300,
        max_depth=6,
        learning_rate=0.05,
        subsample=0.8,
        colsample_bytree=0.8,
        min_child_weight=3,
        reg_alpha=0.0,
        reg_lambda=1.0,
        n_jobs=None,
        random_state=42,
    ):
        self.n_synthetic = n_synthetic
        self.bandwidth = bandwidth
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.learning_rate = learning_rate
        self.subsample = subsample
        self.colsample_bytree = colsample_bytree
        self.min_child_weight = min_child_weight
        self.reg_alpha = reg_alpha
        self.reg_lambda = reg_lambda
        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 _check_is_fitted(self):
        if not hasattr(self, "model_"):
            raise RuntimeError("GeoXGBRegressor is not fitted. Call fit() first.")

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

    def fit(self, X, y, feature_types=None):
        """Fit GeoXGBRegressor.

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

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

        self.X_real_ = X.copy()
        self.y_real_ = y.copy()

        # Fit HVRT on X (y is unused by HVRT — same convention as GeoRFRegressor)
        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 once
        X_real_z = self.hvrt_model_._to_z(X)

        # Single HVRT expansion + kNN y-assignment
        X_syn, y_syn = _make_synthetic_dataset(
            self.hvrt_model_, X_real_z, y,
            n_synthetic=self.n_synthetic,
            seed=self._base_seed,
        )

        # Train XGBoost on real + synthetic
        X_train = np.vstack([X, X_syn])
        y_train = np.concatenate([y, y_syn])

        self.model_ = _XGBRegressor(
            n_estimators=self.n_estimators,
            max_depth=self.max_depth,
            learning_rate=self.learning_rate,
            subsample=self.subsample,
            colsample_bytree=self.colsample_bytree,
            min_child_weight=self.min_child_weight,
            reg_alpha=self.reg_alpha,
            reg_lambda=self.reg_lambda,
            n_jobs=self.n_jobs,
            random_state=self._base_seed,
            verbosity=0,
        )
        self.model_.fit(X_train, y_train)

        return self

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

    def predict(self, X):
        """Predict target values.

        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)
        return self.model_.predict(X)

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

    def feature_importances(self, feature_names=None):
        """XGBoost gain-based feature importances.

        Parameters
        ----------
        feature_names : list[str] | None
            If None, keys are ``'feature_0'``, ``'feature_1'``, ...

        Returns
        -------
        dict[str, float]
            Sorted descending by importance.
        """
        self._check_is_fitted()
        importances = self.model_.feature_importances_
        n = len(importances)
        if feature_names is None:
            feature_names = [f"feature_{i}" for i in range(n)]
        result = dict(zip(feature_names, importances))
        return dict(sorted(result.items(), key=lambda x: x[1], reverse=True))

    def provenance(self):
        """Metadata about the fit.

        Returns
        -------
        dict with keys:
            ``n_real_samples``, ``n_synthetic``, ``n_total_training``,
            ``expansion_ratio``, ``n_estimators``
        """
        self._check_is_fitted()
        n_real = len(self.X_real_)
        return {
            "n_real_samples": n_real,
            "n_synthetic": self.n_synthetic,
            "n_total_training": n_real + self.n_synthetic,
            "expansion_ratio": round(self.n_synthetic / n_real, 3),
            "n_estimators": self.n_estimators,
        }
