"""Public API: SurrogateRandomForestClassifier and SurrogateRandomForestRegressor."""

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from sklearn.utils.multiclass import check_classification_targets
from sklearn.utils.validation import check_is_fitted, validate_data
from joblib import Parallel, delayed

from ._tree import TreeBuilder
from ._validation import validate_sample_weight


def _build_single_tree(
    X, y, sample_weight, tree_params, seed, bootstrap, is_classifier,
):
    """Build a single tree (for parallel execution)."""
    rng = np.random.RandomState(seed)
    n_samples = X.shape[0]

    if bootstrap:
        indices = rng.randint(0, n_samples, size=n_samples)
        oob_mask = np.ones(n_samples, dtype=bool)
        oob_mask[indices] = False
    else:
        indices = np.arange(n_samples)
        oob_mask = np.zeros(n_samples, dtype=bool)

    X_boot = X[indices]
    y_boot = y[indices]
    w_boot = sample_weight[indices]

    builder = TreeBuilder(**tree_params, random_state=seed)
    root = builder.build(X_boot, y_boot, w_boot)

    return root, oob_mask, builder


class SurrogateRandomForestClassifier(ClassifierMixin, BaseEstimator):
    """Random forest classifier with surrogate splits for missing data.

    Combines CART surrogate splits, XGBoost-style NaN direction learning,
    LightGBM histogram binning, and optional CatBoost symmetric trees.
    """

    def __init__(
        self,
        n_estimators=100,
        criterion="gini",
        max_depth=None,
        min_samples_split=2,
        min_samples_leaf=1,
        min_weight_fraction_leaf=0.0,
        max_features="sqrt",
        max_leaf_nodes=None,
        min_impurity_decrease=0.0,
        max_surrogates=5,
        use_histogram=True,
        max_bins=255,
        symmetric_tree=False,
        bootstrap=True,
        oob_score=False,
        n_jobs=None,
        random_state=None,
    ):
        self.n_estimators = n_estimators
        self.criterion = criterion
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.min_weight_fraction_leaf = min_weight_fraction_leaf
        self.max_features = max_features
        self.max_leaf_nodes = max_leaf_nodes
        self.min_impurity_decrease = min_impurity_decrease
        self.max_surrogates = max_surrogates
        self.use_histogram = use_histogram
        self.max_bins = max_bins
        self.symmetric_tree = symmetric_tree
        self.bootstrap = bootstrap
        self.oob_score = oob_score
        self.n_jobs = n_jobs
        self.random_state = random_state

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.input_tags.allow_nan = True
        return tags

    def fit(self, X, y, sample_weight=None):
        X, y = validate_data(self, X, y, accept_sparse=False,
                             ensure_all_finite="allow-nan", dtype=np.float64,
                             multi_output=False)
        sample_weight = validate_sample_weight(sample_weight, X.shape[0])
        check_classification_targets(y)

        # Remove zero-weight samples
        nonzero = sample_weight > 0
        if not nonzero.all():
            X = X[nonzero]
            y = y[nonzero]
            sample_weight = sample_weight[nonzero]

        self.classes_ = np.unique(y)
        self.n_classes_ = len(self.classes_)

        label_map = {c: i for i, c in enumerate(self.classes_)}
        y_mapped = np.array([label_map[v] for v in y], dtype=np.intp)

        rng = np.random.RandomState(self.random_state)
        seeds = rng.randint(0, 2**31, size=self.n_estimators)

        tree_params = dict(
            is_classifier=True,
            n_classes=self.n_classes_,
            criterion=self.criterion,
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            min_weight_fraction_leaf=self.min_weight_fraction_leaf,
            max_features=self.max_features,
            max_leaf_nodes=self.max_leaf_nodes,
            min_impurity_decrease=self.min_impurity_decrease,
            max_surrogates=self.max_surrogates,
            use_histogram=self.use_histogram,
            max_bins=self.max_bins,
            symmetric_tree=self.symmetric_tree,
        )

        results = Parallel(n_jobs=self.n_jobs)(
            delayed(_build_single_tree)(
                X, y_mapped, sample_weight, tree_params, int(seeds[i]), self.bootstrap, True,
            )
            for i in range(self.n_estimators)
        )

        self.estimators_ = []
        self.oob_masks_ = []
        for root, oob_mask, builder in results:
            self.estimators_.append(root)
            self.oob_masks_.append(oob_mask)

        # OOB score
        if self.oob_score and self.bootstrap:
            self.oob_score_ = self._compute_oob_score(X, y_mapped)

        return self

    def predict(self, X):
        check_is_fitted(self)
        proba = self.predict_proba(X)
        return self.classes_[np.argmax(proba, axis=1)]

    def predict_proba(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, accept_sparse=False,
                          ensure_all_finite="allow-nan", dtype=np.float64,
                          reset=False)
        all_proba = np.zeros((X.shape[0], self.n_classes_), dtype=np.float64)

        for tree_root in self.estimators_:
            for i in range(X.shape[0]):
                counts = tree_root.predict_single(X[i])
                total = counts.sum()
                if total > 0:
                    all_proba[i] += counts / total

        all_proba /= len(self.estimators_)
        return all_proba

    def _compute_oob_score(self, X, y):
        n_samples = X.shape[0]
        oob_counts = np.zeros((n_samples, self.n_classes_), dtype=np.float64)
        oob_n = np.zeros(n_samples, dtype=np.float64)

        for tree_root, oob_mask in zip(self.estimators_, self.oob_masks_):
            oob_idx = np.where(oob_mask)[0]
            for i in oob_idx:
                counts = tree_root.predict_single(X[i])
                total = counts.sum()
                if total > 0:
                    oob_counts[i] += counts / total
                    oob_n[i] += 1

        valid = oob_n > 0
        if valid.sum() == 0:
            return 0.0

        oob_pred = np.argmax(oob_counts[valid], axis=1)
        return np.mean(oob_pred == y[valid])


class SurrogateRandomForestRegressor(RegressorMixin, BaseEstimator):
    """Random forest regressor with surrogate splits for missing data."""

    def __init__(
        self,
        n_estimators=100,
        criterion="squared_error",
        max_depth=None,
        min_samples_split=2,
        min_samples_leaf=1,
        min_weight_fraction_leaf=0.0,
        max_features="sqrt",
        max_leaf_nodes=None,
        min_impurity_decrease=0.0,
        max_surrogates=5,
        use_histogram=True,
        max_bins=255,
        symmetric_tree=False,
        bootstrap=True,
        oob_score=False,
        n_jobs=None,
        random_state=None,
    ):
        self.n_estimators = n_estimators
        self.criterion = criterion
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.min_weight_fraction_leaf = min_weight_fraction_leaf
        self.max_features = max_features
        self.max_leaf_nodes = max_leaf_nodes
        self.min_impurity_decrease = min_impurity_decrease
        self.max_surrogates = max_surrogates
        self.use_histogram = use_histogram
        self.max_bins = max_bins
        self.symmetric_tree = symmetric_tree
        self.bootstrap = bootstrap
        self.oob_score = oob_score
        self.n_jobs = n_jobs
        self.random_state = random_state

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.input_tags.allow_nan = True
        return tags

    def fit(self, X, y, sample_weight=None):
        X, y = validate_data(self, X, y, accept_sparse=False,
                             ensure_all_finite="allow-nan", dtype=np.float64,
                             multi_output=False)
        sample_weight = validate_sample_weight(sample_weight, X.shape[0])
        y = y.astype(np.float64)

        # Remove zero-weight samples
        nonzero = sample_weight > 0
        if not nonzero.all():
            X = X[nonzero]
            y = y[nonzero]
            sample_weight = sample_weight[nonzero]

        rng = np.random.RandomState(self.random_state)
        seeds = rng.randint(0, 2**31, size=self.n_estimators)

        tree_params = dict(
            is_classifier=False,
            criterion=self.criterion,
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            min_weight_fraction_leaf=self.min_weight_fraction_leaf,
            max_features=self.max_features,
            max_leaf_nodes=self.max_leaf_nodes,
            min_impurity_decrease=self.min_impurity_decrease,
            max_surrogates=self.max_surrogates,
            use_histogram=self.use_histogram,
            max_bins=self.max_bins,
            symmetric_tree=self.symmetric_tree,
        )

        results = Parallel(n_jobs=self.n_jobs)(
            delayed(_build_single_tree)(
                X, y, sample_weight, tree_params, int(seeds[i]), self.bootstrap, False,
            )
            for i in range(self.n_estimators)
        )

        self.estimators_ = []
        self.oob_masks_ = []
        for root, oob_mask, builder in results:
            self.estimators_.append(root)
            self.oob_masks_.append(oob_mask)

        if self.oob_score and self.bootstrap:
            self.oob_score_ = self._compute_oob_score(X, y)

        return self

    def predict(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, accept_sparse=False,
                          ensure_all_finite="allow-nan", dtype=np.float64,
                          reset=False)
        predictions = np.zeros(X.shape[0], dtype=np.float64)

        for tree_root in self.estimators_:
            for i in range(X.shape[0]):
                val = tree_root.predict_single(X[i])
                predictions[i] += val[0]  # mean value

        predictions /= len(self.estimators_)
        return predictions

    def _compute_oob_score(self, X, y):
        n_samples = X.shape[0]
        oob_sum = np.zeros(n_samples, dtype=np.float64)
        oob_n = np.zeros(n_samples, dtype=np.float64)

        for tree_root, oob_mask in zip(self.estimators_, self.oob_masks_):
            oob_idx = np.where(oob_mask)[0]
            for i in oob_idx:
                val = tree_root.predict_single(X[i])
                oob_sum[i] += val[0]
                oob_n[i] += 1

        valid = oob_n > 0
        if valid.sum() == 0:
            return 0.0

        oob_pred = oob_sum[valid] / oob_n[valid]
        ss_res = np.sum((y[valid] - oob_pred) ** 2)
        ss_tot = np.sum((y[valid] - y[valid].mean()) ** 2)
        return 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0
