"""GeoRFClassifier: geometry-aware random forest for classification."""
import numpy as np
from joblib import Parallel, delayed
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import LabelEncoder

from hvrt import HVRT

from ._base import _GeoRFBase
from ._utils import fit_one_tree_clf


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

    Each tree is trained on a completely fresh synthetic dataset drawn from
    per-class HVRT density models.  No tree ever sees a real sample; no two
    trees share a training point.

    Parameters
    ----------
    n_estimators : int, default=100
        Number of trees in the ensemble.
    n_samples_per_tree : int, default=500
        Synthetic samples generated (in total, across classes) 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 strategy per split.  ``None`` uses all features;
        generative diversity alone is sufficient.  ``'sqrt'``, ``'log2'``,
        int and float are also accepted (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.  ``None`` = 1 core, ``-1`` = all cores.
    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 classifier.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training features.  Categorical features must be pre-encoded to
            numeric values.
        y : array-like of shape (n_samples,)
            Target class labels (numeric or string).
        feature_types : list[str] | None
            Per-column type: ``'continuous'`` or ``'categorical'``.
            ``None`` treats all columns as continuous.

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

        # Encode labels to integers 0..n_classes-1
        self.le_ = LabelEncoder()
        y_encoded = self.le_.fit_transform(y)
        self.classes_ = self.le_.classes_

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

        # ---- fit one HVRT per class ----
        unique_encoded = np.unique(y_encoded)
        self.hvrt_models_ = {}
        self.class_counts_ = {}

        for enc_label in unique_encoded:
            mask = y_encoded == enc_label
            X_class = X[mask]
            self.class_counts_[int(enc_label)] = len(X_class)
            hvrt = HVRT(bandwidth=self.bandwidth, random_state=self._base_seed)
            hvrt.fit(X_class, feature_types=feature_types)
            self.hvrt_models_[int(enc_label)] = hvrt

        n_total = len(X)
        class_encoded_labels = [int(c) for c in unique_encoded]

        # Compute synthetic class sizes (proportional, floor at 1)
        n_per_class_list = [
            max(1, round(self.n_samples_per_tree * self.class_counts_[c] / n_total))
            for c in class_encoded_labels
        ]

        # ---- fit trees in parallel ----
        # Each worker expands synthetic data + fits its tree with a unique seed.
        # random_state=seed is passed to both the HVRT copy (KDE sampling) and
        # the DecisionTree (tie-breaking) to guarantee full determinism.
        seeds = [self._base_seed + i for i in range(self.n_estimators)]

        self.trees_ = Parallel(n_jobs=self.n_jobs)(
            delayed(fit_one_tree_clf)(
                seed, self.hvrt_models_, class_encoded_labels,
                n_per_class_list, self._tree_kwargs(),
            )
            for seed in seeds
        )

        return self

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

    def predict_proba(self, X):
        """Average class probabilities across all trees.

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

        Returns
        -------
        ndarray of shape (n_samples, n_classes)
            Each row sums to 1.0.
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)

        # Stack (n_trees, n_samples, n_classes) then average over trees
        all_probas = np.array([tree.predict_proba(X) for tree in self.trees_])
        return all_probas.mean(axis=0)

    def predict(self, X):
        """Predict class labels.

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

        Returns
        -------
        array of shape (n_samples,)
            Class labels in the original label type.
        """
        proba = self.predict_proba(X)
        pred_encoded = np.argmax(proba, axis=1)
        return self.le_.inverse_transform(pred_encoded)

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

    def _score_tree(self, tree):
        """AUC for a single tree on the real seed data."""
        proba = tree.predict_proba(self.X_real_)
        if len(self.classes_) == 2:
            return roc_auc_score(self.y_encoded_, proba[:, 1])
        return roc_auc_score(
            self.y_encoded_, proba, multi_class="ovr", average="macro"
        )

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

        Disagreement between two trees is the fraction of real-data samples
        on which they predict different class labels.

        Returns
        -------
        float in [0, 1]
        """
        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

        total = 0.0
        count = 0
        for i in range(n_trees):
            for j in range(i + 1, n_trees):
                total += np.mean(preds[i] != preds[j])
                count += 1

        return float(total / count)

    def predict_proba_std(self, X):
        """Per-sample std of class probabilities across trees.

        Measures how much the ensemble disagrees on each sample's class
        probabilities — a calibration-free uncertainty estimate.

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

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

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