import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

from ._xgboost import xgb_opt_classifier, xgb_opt_regressor


class LCETreeClassifier(ClassifierMixin, BaseEstimator):
    """
    A LCE Tree classifier.


    Parameters
    ----------
    n_classes_in : int, default=None
        The number of classes from the input data.

    criterion : {"gini", "entropy"}, default="gini"
        The function to measure the quality of a split. Supported criteria are
        "gini" for the Gini impurity and "entropy" for the information gain.

    splitter : {"best", "random"}, default="best"
        The strategy used to choose the split at each node. Supported strategies
        are "best" to choose the best split and "random" to choose the best random
        split.

    max_depth : int, default=2
        The maximum depth of a tree.

    max_features : int, float or {"auto", "sqrt", "log"}, default=None
        The number of features to consider when looking for the best split:

        - If int, then consider `max_features` features at each split.
        - If float, then `max_features` is a fraction and
          `round(max_features * n_features)` features are considered at each
          split.
        - If "auto", then `max_features=sqrt(n_features)`.
        - If "sqrt", then `max_features=sqrt(n_features)` (same as "auto").
        - If "log2", then `max_features=log2(n_features)`.
        - If None, then `max_features=n_features`.

        Note: the search for a split does not stop until at least one
        valid partition of the node samples is found, even if it requires to
        effectively inspect more than ``max_features`` features.

    min_samples_leaf : int or float, default=5
        The minimum number of samples required to be at a leaf node.
        A split point at any depth will only be considered if it leaves at
        least ``min_samples_leaf`` training samples in each of the left and
        right branches.

        - If int, then consider `min_samples_leaf` as the minimum number.
        - If float, then `min_samples_leaf` is a fraction and
          `ceil(min_samples_leaf * n_samples)` are the minimum
          number of samples for each node.

    n_iter: int, default=10
        Number of iterations to set the hyperparameters of each node base
        classifier in Hyperopt.

    metric: string, default="accuracy"
        The score of the base classifier (XGBoost) optimized by Hyperopt. Supported metrics
        are the ones from `scikit-learn <https://scikit-learn.org/stable/modules/model_evaluation.html>`_.

    xgb_max_n_estimators : int, default=100
        The maximum number of XGBoost estimators. The number of estimators of
        XGBoost corresponds to the number of boosting rounds.

    xgb_n_estimators_step : int, default=10
        Spacing between XGBoost n_estimators. The range of XGBoost n_estimators
        for hyperparameter optimization (Hyperopt) is:
        range(1, xgb_max_n_estimators+xgb_n_estimators_step, xgb_n_estimators_step).

    xgb_max_depth : int, default=10
        Maximum tree depth for XGBoost base learners. The range of XGBoost max_depth
        for Hyperopt is: range(1, xgb_max_depth+1).

    xgb_min_learning_rate : float, default=0.05
        Minimum learning rate of XGBoost. The learning rate corresponds to the
        step size shrinkage used in update to prevent overfitting. After each
        boosting step, we can directly get the weights of new features,
        and the learning rate shrinks the feature weights to make the boosting
        process more conservative.

    xgb_max_learning_rate : float, default=0.5
        Maximum learning rate of XGBoost.

    xgb_learning_rate_step : float, default=0.05
        Spacing between XGBoost learning_rate. The range of XGBoost learning_rate
        for hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_learning_rate, xgb_max_learning_rate+xgb_learning_rate_step, xgb_learning_rate_step).

    xgb_booster : {"dart", "gblinear", "gbtree"}, default="gbtree"
        The type of booster to use. "gbtree" and "dart" use tree based models
        while "gblinear" uses linear functions.

    xgb_min_gamma : float, default=0.05
        Minimum gamma of XGBoost. Gamma corresponds to the minimum loss reduction
        required to make a further partition on a leaf node of the tree.
        The larger gamma is, the more conservative XGBoost algorithm will be.

    xgb_max_gamma : float, default=0.5
        Maximum gamma of XGBoost.

    xgb_gamma_step : float, default=0.05,
        Spacing between XGBoost gamma. The range of XGBoost gamma for hyperparameter
        optimization (Hyperopt) is:
        np.arange(xgb_min_gamma, xgb_max_gamma+xgb_gamma_step, xgb_gamma_step).

    xgb_min_min_child_weight : int, default=3
        Minimum min_child_weight of XGBoost. min_child_weight defines the
        minimum sum of instance weight (hessian) needed in a child. If the tree
        partition step results in a leaf node with the sum of instance weight
        less than min_child_weight, then the building process will give up further
        partitioning. The larger min_child_weight is, the more conservative XGBoost
        algorithm will be.

    xgb_max_min_child_weight : int, default=10
        Minimum min_child_weight of XGBoost.

    xgb_min_child_weight_step : int, default=1,
        Spacing between XGBoost min_child_weight. The range of XGBoost min_child_weight
        for hyperparameter optimization (Hyperopt) is:
        range(xgb_min_min_child_weight, xgb_max_min_child_weight+xgb_min_child_weight_step, xgb_min_child_weight_step).

    xgb_subsample : float, default=0.8
        XGBoost subsample ratio of the training instances. Setting it to 0.5 means
        that XGBoost would randomly sample half of the training data prior to
        growing trees. and this will prevent overfitting. Subsampling will occur
        once in every boosting iteration.

    xgb_colsample_bytree : float, default=0.8
        XGBoost subsample ratio of columns when constructing each tree.
        Subsampling occurs once for every tree constructed.

    xgb_colsample_bylevel : float, default=1.0
        XGBoost subsample ratio of columns for each level. Subsampling occurs
        once for every new depth level reached in a tree. Columns are subsampled
        from the set of columns chosen for the current tree.

    xgb_colsample_bynode : float, default=1.0
        XGBoost subsample ratio of columns for each node (split). Subsampling
        occurs once every time a new split is evaluated. Columns are subsampled
        from the set of columns chosen for the current level.

    xgb_min_reg_alpha : float, default=0.01
        Minimum reg_alpha of XGBoost. reg_alpha corresponds to the L1 regularization
        term on the weights. Increasing this value will make XGBoost model more
        conservative.

    xgb_max_reg_alpha : float, default=0.1
        Maximum reg_alpha of XGBoost.

    xgb_reg_alpha_step : float, default=0.05
        Spacing between XGBoost reg_alpha. The range of XGBoost reg_alpha for
        hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_reg_alpha, xgb_max_reg_alpha+xgb_reg_alpha_step, xgb_reg_alpha_step).

    xgb_min_reg_lambda : float, default=0.01
        Minimum reg_lambda of XGBoost. reg_lambda corresponds to the L2 regularization
        term on the weights. Increasing this value will make XGBoost model more
        conservative.

    xgb_max_reg_lambda : float, default=0.1
        Maximum reg_lambda of XGBoost.

    xgb_reg_lambda_step : float, default=0.05
        Spacing between XGBoost reg_lambda. The range of XGBoost reg_lambda for
        hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_reg_lambda, xgb_max_reg_lambda+xgb_reg_lambda_step, xgb_reg_lambda_step).

    n_jobs : int, default=None
        The number of jobs to run in parallel.
        ``None`` means 1. ``-1`` means using all processors.

    random_state : int, RandomState instance or None, default=None
        Controls the randomness of the sampling of the features to consider when
        looking for the best split at each node (if ``max_features < n_features``),
        the base classifier (XGBoost) and the Hyperopt algorithm.

    verbose : int, default=0
        Controls the verbosity when fitting.

    Attributes
    ----------
    classes_ : ndarray of shape (n_classes,) or a list of such arrays
        The classes labels.

    n_features_in_ : int
        The number of features when ``fit`` is performed.
    """

    def __init__(
        self,
        n_classes_in=None,
        criterion="gini",
        splitter="best",
        max_depth=2,
        max_features=None,
        min_samples_leaf=5,
        n_iter=10,
        metric="accuracy",
        xgb_max_n_estimators=100,
        xgb_n_estimators_step=10,
        xgb_max_depth=10,
        xgb_min_learning_rate=0.05,
        xgb_max_learning_rate=0.5,
        xgb_learning_rate_step=0.05,
        xgb_booster="gbtree",
        xgb_min_gamma=0.05,
        xgb_max_gamma=0.5,
        xgb_gamma_step=0.05,
        xgb_min_min_child_weight=3,
        xgb_max_min_child_weight=10,
        xgb_min_child_weight_step=1,
        xgb_subsample=0.8,
        xgb_colsample_bytree=0.8,
        xgb_colsample_bylevel=1.0,
        xgb_colsample_bynode=1.0,
        xgb_min_reg_alpha=0.01,
        xgb_max_reg_alpha=0.1,
        xgb_reg_alpha_step=0.05,
        xgb_min_reg_lambda=0.01,
        xgb_max_reg_lambda=0.1,
        xgb_reg_lambda_step=0.05,
        n_jobs=None,
        random_state=None,
        verbose=0,
    ):
        self.n_classes_in = n_classes_in
        self.criterion = criterion
        self.splitter = splitter
        self.max_depth = max_depth
        self.max_features = max_features
        self.min_samples_leaf = min_samples_leaf
        self.n_iter = n_iter
        self.metric = metric
        self.xgb_max_n_estimators = xgb_max_n_estimators
        self.xgb_n_estimators_step = xgb_n_estimators_step
        self.xgb_max_depth = xgb_max_depth
        self.xgb_min_learning_rate = xgb_min_learning_rate
        self.xgb_max_learning_rate = xgb_max_learning_rate
        self.xgb_learning_rate_step = xgb_learning_rate_step
        self.xgb_booster = xgb_booster
        self.xgb_min_gamma = xgb_min_gamma
        self.xgb_max_gamma = xgb_max_gamma
        self.xgb_gamma_step = xgb_gamma_step
        self.xgb_min_min_child_weight = xgb_min_min_child_weight
        self.xgb_max_min_child_weight = xgb_max_min_child_weight
        self.xgb_min_child_weight_step = xgb_min_child_weight_step
        self.xgb_subsample = xgb_subsample
        self.xgb_colsample_bytree = xgb_colsample_bytree
        self.xgb_colsample_bylevel = xgb_colsample_bylevel
        self.xgb_colsample_bynode = xgb_colsample_bynode
        self.xgb_min_reg_alpha = xgb_min_reg_alpha
        self.xgb_max_reg_alpha = xgb_max_reg_alpha
        self.xgb_reg_alpha_step = xgb_reg_alpha_step
        self.xgb_min_reg_lambda = xgb_min_reg_lambda
        self.xgb_max_reg_lambda = xgb_max_reg_lambda
        self.xgb_reg_lambda_step = xgb_reg_lambda_step
        self.n_jobs = n_jobs
        self.random_state = random_state
        self.verbose = verbose

    def fit(self, X, y):
        """
        Build a LCE tree from the training set (X, y).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The training input samples.

        y : array-like of shape (n_samples,)
            The class labels.

        Returns
        -------
        self : object
        """
        self.classes_ = np.unique(y)
        self.n_features_in_ = X.shape[1]

        def _build_tree(X, y):
            """Build a LCE tree."""
            global index_node_global

            def _create_node(X, y, depth, container):
                """Create a node in the tree."""
                # Add XGBoost predictions as features to the dataset
                model_node = xgb_opt_classifier(
                    X,
                    y,
                    n_iter=self.n_iter,
                    metric=self.metric,
                    n_estimators=self.xgb_max_n_estimators,
                    n_estimators_step=self.xgb_n_estimators_step,
                    max_depth=self.xgb_max_depth,
                    min_learning_rate=self.xgb_min_learning_rate,
                    max_learning_rate=self.xgb_max_learning_rate,
                    learning_rate_step=self.xgb_learning_rate_step,
                    booster=self.xgb_booster,
                    min_gamma=self.xgb_min_gamma,
                    max_gamma=self.xgb_max_gamma,
                    gamma_step=self.xgb_gamma_step,
                    min_min_child_weight=self.xgb_min_min_child_weight,
                    max_min_child_weight=self.xgb_max_min_child_weight,
                    min_child_weight_step=self.xgb_min_child_weight_step,
                    subsample=self.xgb_subsample,
                    colsample_bytree=self.xgb_colsample_bytree,
                    colsample_bylevel=self.xgb_colsample_bylevel,
                    colsample_bynode=self.xgb_colsample_bynode,
                    min_reg_alpha=self.xgb_min_reg_alpha,
                    max_reg_alpha=self.xgb_max_reg_alpha,
                    reg_alpha_step=self.xgb_reg_alpha_step,
                    min_reg_lambda=self.xgb_min_reg_lambda,
                    max_reg_lambda=self.xgb_max_reg_lambda,
                    reg_lambda_step=self.xgb_reg_lambda_step,
                    n_jobs=self.n_jobs,
                    random_state=self.random_state,
                )
                pred_proba = np.around(model_node.predict_proba(X), 6)

                c = 0
                for i in range(0, self.n_classes_in):
                    X = np.insert(X, X.shape[1], 0, axis=1)
                    if i in y:
                        if np.unique(y).size == 1:
                            X[:, -1] = pred_proba[:, 1]
                        else:
                            X[:, -1] = pred_proba[:, c]
                            c += 1

                # Missing data information
                num_nans = np.isnan(X).any(axis=1).sum()
                if num_nans > 0:
                    missing = True
                    if num_nans == y.size:
                        missing_only = True
                    else:
                        missing_only = False
                else:
                    missing = False
                    missing_only = False

                # Split
                split_val_conditions = [y.size > 1, missing_only == False]
                if all(split_val_conditions):
                    split = DecisionTreeClassifier(
                        criterion=self.criterion,
                        splitter=self.splitter,
                        max_depth=1,
                        max_features=self.max_features,
                        random_state=self.random_state,
                    )
                    if missing:
                        nans = np.isnan(X).any(axis=1)
                        split.fit(X[~nans], y[~nans])
                    else:
                        split.fit(X, y)
                else:
                    split = None

                # Node information
                node = {
                    "index": container["index_node_global"],
                    "model": model_node,
                    "data": (X, y),
                    "classes_in": np.unique(y),
                    "num_classes": self.n_classes_in,
                    "split": split,
                    "missing": {"missing": missing, "missing_only": missing_only},
                    "missing_side": None,
                    "children": {"left": None, "right": None},
                    "depth": depth,
                }
                container["index_node_global"] += 1
                return node

            def _splitter(node):
                """Perform the split of a node."""
                # Extract data
                X, y = node["data"]
                depth = node["depth"]
                split = node["split"]
                missing = node["missing"]["missing"]
                missing_only = node["missing"]["missing_only"]

                did_split = False
                data = None

                # Perform split if the conditions are met
                stopping_criteria = [
                    depth >= 0,
                    depth < self.max_depth,
                    np.unique(y).size > 1,
                    missing_only == False,
                ]

                if all(stopping_criteria):
                    if missing:
                        nans = np.isnan(X).any(axis=1)
                        X_withoutnans, y_withoutnans = X[~nans], y[~nans]
                        leafs = split.apply(X_withoutnans)
                        (X_left, y_left), (X_right, y_right) = (
                            np.squeeze(X_withoutnans[np.argwhere(leafs == 1), :]),
                            np.squeeze(y_withoutnans[np.argwhere(leafs == 1)]),
                        ), (
                            np.squeeze(X_withoutnans[np.argwhere(leafs == 2), :]),
                            np.squeeze(y_withoutnans[np.argwhere(leafs == 2)]),
                        )
                    else:
                        leafs = split.apply(X)
                        (X_left, y_left), (X_right, y_right) = (
                            np.squeeze(X[np.argwhere(leafs == 1), :]),
                            np.squeeze(y[np.argwhere(leafs == 1)]),
                        ), (
                            np.squeeze(X[np.argwhere(leafs == 2), :]),
                            np.squeeze(y[np.argwhere(leafs == 2)]),
                        )

                    N_left, N_right = y_left.size, y_right.size

                    split_conditions = [
                        N_left >= self.min_samples_leaf,
                        N_right >= self.min_samples_leaf,
                    ]

                    if all(split_conditions):
                        did_split = True

                        if N_left == 1:
                            X_left = X_left.reshape(-1, 1).T
                            node["missing_side"] = "left"
                            if missing:
                                X_left = np.append(X_left, X[nans], axis=0)
                                y_left = np.append([y_left], y[nans], axis=0)

                        if N_right == 1:
                            X_right = X_right.reshape(-1, 1).T
                            if N_left > 1:
                                node["missing_side"] = "right"
                                if missing:
                                    X_right = np.append(X_right, X[nans], axis=0)
                                    y_right = np.append([y_right], y[nans], axis=0)

                        score_conditions = [N_left > 1, N_right > 1]
                        if all(score_conditions):
                            if split.score(X_left, y_left) > split.score(
                                X_right, y_right
                            ):
                                node["missing_side"] = "left"
                                if missing:
                                    X_left = np.append(X_left, X[nans], axis=0)
                                    y_left = np.append(y_left, y[nans], axis=0)
                            else:
                                node["missing_side"] = "right"
                                if missing:
                                    X_right = np.append(X_right, X[nans], axis=0)
                                    y_right = np.append(y_right, y[nans], axis=0)

                        data = [(X_left, y_left), (X_right, y_right)]

                result = {"did_split": did_split, "data": data}
                return result

            def _split_traverse_node(node, container):
                """Process splitting results and continue with child nodes."""
                # Perform split and collect result
                result = _splitter(node)

                # Return terminal node if no split
                if not result["did_split"]:
                    if self.verbose > 0 and self.n_jobs == None:
                        depth_spacing_str = " ".join([" "] * node["depth"])
                        print(
                            " {}*leaf {} @ depth {}: Unique_y {},  N_samples {}".format(
                                depth_spacing_str,
                                node["index"],
                                node["depth"],
                                np.unique(node["data"][1]),
                                np.unique(node["data"][1], return_counts=True)[1],
                            )
                        )
                    return
                del node["data"]

                # Extract splitting results
                (X_left, y_left), (X_right, y_right) = result["data"]

                # Report created node to user
                if self.verbose > 0 and self.n_jobs == None:
                    depth_spacing_str = " ".join([" "] * node["depth"])
                    print(
                        " {}node {} @ depth {}: dataset={}, N_left={}, N_right={}".format(
                            depth_spacing_str,
                            node["index"],
                            node["depth"],
                            (X_left.shape[0] + X_right.shape[0], X_left.shape[1]),
                            X_left.shape[0],
                            X_right.shape[0],
                        )
                    )

                # Create child nodes
                node["children"]["left"] = _create_node(
                    X_left, y_left, node["depth"] + 1, container
                )
                node["children"]["right"] = _create_node(
                    X_right, y_right, node["depth"] + 1, container
                )

                # Split nodes
                _split_traverse_node(node["children"]["left"], container)
                _split_traverse_node(node["children"]["right"], container)

            container = {"index_node_global": 0}
            if self.verbose > 0 and self.n_jobs == None:
                print("\nNew Tree")
            root = _create_node(X, y, 0, container)
            _split_traverse_node(root, container)
            return root

        self.tree = _build_tree(X, y)
        return self

    def predict_proba(self, X):
        """
        Predict class probabilities for X.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The training input samples.

        Returns
        -------
        y : ndarray of shape (n_samples,)
            The class probabilities of the input samples.
        """

        def _base_proba(node, X):
            y_pred = np.around(node["model"].predict_proba(X[:, 1:]), 6)
            d = 0
            for j in range(0, node["num_classes"]):
                X = np.insert(X, X.shape[1], 0, axis=1)
                if j in node["classes_in"]:
                    if node["classes_in"].size == 1:
                        X[:, -1] = y_pred[:, 1]
                    else:
                        X[:, -1] = y_pred[:, d]
                        d += 1
            return X

        def _predict_proba(node, X, y_pred_final=None):
            if X.ndim == 1:
                X = _base_proba(node, X.reshape(-1, 1).T)
            else:
                X = _base_proba(node, X)

            no_children = (
                node["children"]["left"] is None and node["children"]["right"] is None
            )
            if no_children:
                y_pred = np.column_stack((X[:, :1], X[:, -node["num_classes"] :]))
                if y_pred_final is not None:
                    y_pred_final = np.concatenate((y_pred_final, y_pred), axis=0)
                else:
                    y_pred_final = y_pred
                return y_pred_final

            else:
                if np.isnan(X).sum() > 0:
                    nans = np.isnan(X).any(axis=1)
                    leafs = node["split"].apply(X[~nans, 1:])
                    X_left, X_right = np.squeeze(
                        X[~nans][np.argwhere(leafs == 1), :]
                    ), np.squeeze(X[~nans][np.argwhere(leafs == 2), :])
                    if node["missing_side"] == "left":
                        X_left, X_right = (
                            np.concatenate((X_left, X[nans]), axis=0),
                            X_right,
                        )
                    else:
                        X_left, X_right = X_left, np.concatenate(
                            (X_right, X[nans]), axis=0
                        )
                else:
                    leafs = node["split"].apply(X[:, 1:])
                    X_left, X_right = np.squeeze(
                        X[np.argwhere(leafs == 1), :]
                    ), np.squeeze(X[np.argwhere(leafs == 2), :])

                if len(X_left) > 0:
                    y_pred_final = _predict_proba(
                        node["children"]["left"], X_left, y_pred_final
                    )
                if len(X_right) > 0:
                    y_pred_final = _predict_proba(
                        node["children"]["right"], X_right, y_pred_final
                    )
                return y_pred_final

        index = np.arange(0, X.shape[0]).reshape(-1, 1)
        X = np.concatenate((index, X), axis=1)
        y_pred = _predict_proba(self.tree, X, None)
        y_pred = y_pred[y_pred[:, 0].argsort()]
        y_pred = y_pred[:, 1:]
        return y_pred

    def set_params(self, **params):
        """
        Set the parameters of the estimator.

        Parameters
        ----------
        **params : dict
            Estimator parameters.

        Returns
        -------
        self : object
        """
        if not params:
            return self

        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                self.kwargs[key] = value
        return self


class LCETreeRegressor(RegressorMixin, BaseEstimator):
    """
    A LCE Tree regressor.


    Parameters
    ----------
    criterion : {"squared_error", "friedman_mse", "absolute_error", "poisson"}, default="squared_error"
        The function to measure the quality of a split. Supported criteria are "squared_error" for
        the mean squared error, which is equal to variance reduction as feature selection
        criterion and minimizes the L2 loss using the mean of each terminal node,
        "friedman_mse", which uses mean squared error with Friedman's improvement score
        for potential splits, "absolute_error" for the mean absolute error, which
        minimizes the L1 loss using the median of each terminal node, and "poisson"
        which uses reduction in Poisson deviance to find splits.

    splitter : {"best", "random"}, default="best"
        The strategy used to choose the split at each node. Supported strategies
        are "best" to choose the best split and "random" to choose the best random
        split.

    max_depth : int, default=2
        The maximum depth of a tree.

    max_features : int, float or {"auto", "sqrt", "log"}, default=None
        The number of features to consider when looking for the best split:

        - If int, then consider `max_features` features at each split.
        - If float, then `max_features` is a fraction and
          `round(max_features * n_features)` features are considered at each
          split.
        - If "auto", then `max_features=sqrt(n_features)`.
        - If "sqrt", then `max_features=sqrt(n_features)` (same as "auto").
        - If "log2", then `max_features=log2(n_features)`.
        - If None, then `max_features=n_features`.

        Note: the search for a split does not stop until at least one
        valid partition of the node samples is found, even if it requires to
        effectively inspect more than ``max_features`` features.

    min_samples_leaf : int or float, default=5
        The minimum number of samples required to be at a leaf node.
        A split point at any depth will only be considered if it leaves at
        least ``min_samples_leaf`` training samples in each of the left and
        right branches.

        - If int, then consider `min_samples_leaf` as the minimum number.
        - If float, then `min_samples_leaf` is a fraction and
          `ceil(min_samples_leaf * n_samples)` are the minimum
          number of samples for each node.

    n_iter: int, default=10
        Number of iterations to set the hyperparameters of each node base
        regressor (XGBoost) in Hyperopt.

    metric: string, default="neg_mean_squared_error"
        The score of the base regressor (XGBoost) optimized by Hyperopt. Supported metrics
        are the ones from `scikit-learn <https://scikit-learn.org/stable/modules/model_evaluation.html>`_.

    xgb_max_n_estimators : int, default=100
        The maximum number of XGBoost estimators. The number of estimators of
        XGBoost corresponds to the number of boosting rounds.

    xgb_n_estimators_step : int, default=10
        Spacing between XGBoost n_estimators. The range of XGBoost n_estimators
        for hyperparameter optimization (Hyperopt) is:
        range(1, xgb_max_n_estimators+xgb_n_estimators_step, xgb_n_estimators_step).

    xgb_max_depth : int, default= 10
        Maximum tree depth for XGBoost base learners. The range of XGBoost max_depth
        for Hyperopt is: range(1, xgb_max_depth+1).

    xgb_min_learning_rate : float, default=0.05
        Minimum learning rate of XGBoost. The learning rate corresponds to the
        step size shrinkage used in update to prevent overfitting. After each
        boosting step, we can directly get the weights of new features,
        and the learning rate shrinks the feature weights to make the boosting
        process more conservative.

    xgb_max_learning_rate : float, default=0.5
        Maximum learning rate of XGBoost.

    xgb_learning_rate_step : float, default=0.05
        Spacing between XGBoost learning_rate. The range of XGBoost learning_rate
        for hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_learning_rate, xgb_max_learning_rate+xgb_learning_rate_step, xgb_learning_rate_step).

    xgb_booster : {"dart", "gblinear", "gbtree"}, default="gbtree"
        The type of booster to use. "gbtree" and "dart" use tree based models
        while "gblinear" uses linear functions.

    xgb_min_gamma : float, default=0.05
        Minimum gamma of XGBoost. Gamma corresponds to the minimum loss reduction
        required to make a further partition on a leaf node of the tree.
        The larger gamma is, the more conservative XGBoost algorithm will be.

    xgb_max_gamma : float, default=0.5
        Maximum gamma of XGBoost.

    xgb_gamma_step : float, default=0.05,
        Spacing between XGBoost gamma. The range of XGBoost gamma for hyperparameter
        optimization (Hyperopt) is:
        np.arange(xgb_min_gamma, xgb_max_gamma+xgb_gamma_step, xgb_gamma_step).

    xgb_min_min_child_weight : int, default=3
        Minimum min_child_weight of XGBoost. min_child_weight defines the
        minimum sum of instance weight (hessian) needed in a child. If the tree
        partition step results in a leaf node with the sum of instance weight
        less than min_child_weight, then the building process will give up further
        partitioning. The larger min_child_weight is, the more conservative XGBoost
        algorithm will be.

    xgb_max_min_child_weight : int, default=10
        Minimum min_child_weight of XGBoost.

    xgb_min_child_weight_step : int, default=1,
        Spacing between XGBoost min_child_weight. The range of XGBoost min_child_weight
        for hyperparameter optimization (Hyperopt) is:
        range(xgb_min_min_child_weight, xgb_max_min_child_weight+xgb_min_child_weight_step, xgb_min_child_weight_step).

    xgb_subsample : float, default=0.8
        XGBoost subsample ratio of the training instances. Setting it to 0.5 means
        that XGBoost would randomly sample half of the training data prior to
        growing trees. and this will prevent overfitting. Subsampling will occur
        once in every boosting iteration.

    xgb_colsample_bytree : float, default=0.8
        XGBoost subsample ratio of columns when constructing each tree.
        Subsampling occurs once for every tree constructed.

    xgb_colsample_bylevel : float, default=1.0
        XGBoost subsample ratio of columns for each level. Subsampling occurs
        once for every new depth level reached in a tree. Columns are subsampled
        from the set of columns chosen for the current tree.

    xgb_colsample_bynode : float, default=1.0
        XGBoost subsample ratio of columns for each node (split). Subsampling
        occurs once every time a new split is evaluated. Columns are subsampled
        from the set of columns chosen for the current level.

    xgb_min_reg_alpha : float, default=0.01
        Minimum reg_alpha of XGBoost. reg_alpha corresponds to the L1 regularization
        term on the weights. Increasing this value will make XGBoost model more
        conservative.

    xgb_max_reg_alpha : float, default=0.1
        Maximum reg_alpha of XGBoost.

    xgb_reg_alpha_step : float, default=0.05
        Spacing between XGBoost reg_alpha. The range of XGBoost reg_alpha for
        hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_reg_alpha, xgb_max_reg_alpha+xgb_reg_alpha_step, xgb_reg_alpha_step).

    xgb_min_reg_lambda : float, default=0.01
        Minimum reg_lambda of XGBoost. reg_lambda corresponds to the L2 regularization
        term on the weights. Increasing this value will make XGBoost model more
        conservative.

    xgb_max_reg_lambda : float, default=0.1
        Maximum reg_lambda of XGBoost.

    xgb_reg_lambda_step : float, default=0.05
        Spacing between XGBoost reg_lambda. The range of XGBoost reg_lambda for
        hyperparameter optimization (Hyperopt) is:
        np.arange(xgb_min_reg_lambda, xgb_max_reg_lambda+xgb_reg_lambda_step, xgb_reg_lambda_step).

    n_jobs : int, default=None
        The number of jobs to run in parallel.
        ``None`` means 1. ``-1`` means using all processors.

    random_state : int, RandomState instance or None, default=None
        Controls the randomness of the sampling of the features to consider when
        looking for the best split at each node (if ``max_features < n_features``),
        the base classifier (XGBoost) and the Hyperopt algorithm.

    verbose : int, default=0
        Controls the verbosity when fitting.

    Attributes
    ----------
    n_features_in_ : int
        The number of features when ``fit`` is performed.
    """

    def __init__(
        self,
        criterion="squared_error",
        splitter="best",
        max_depth=2,
        max_features=None,
        min_samples_leaf=5,
        n_iter=10,
        metric="neg_mean_squared_error",
        xgb_max_n_estimators=100,
        xgb_n_estimators_step=10,
        xgb_max_depth=10,
        xgb_min_learning_rate=0.05,
        xgb_max_learning_rate=0.5,
        xgb_learning_rate_step=0.05,
        xgb_booster="gbtree",
        xgb_min_gamma=0.05,
        xgb_max_gamma=0.5,
        xgb_gamma_step=0.05,
        xgb_min_min_child_weight=3,
        xgb_max_min_child_weight=10,
        xgb_min_child_weight_step=1,
        xgb_subsample=0.8,
        xgb_colsample_bytree=0.8,
        xgb_colsample_bylevel=1.0,
        xgb_colsample_bynode=1.0,
        xgb_min_reg_alpha=0.01,
        xgb_max_reg_alpha=0.1,
        xgb_reg_alpha_step=0.05,
        xgb_min_reg_lambda=0.01,
        xgb_max_reg_lambda=0.1,
        xgb_reg_lambda_step=0.05,
        n_jobs=None,
        random_state=None,
        verbose=0,
    ):
        self.criterion = criterion
        self.splitter = splitter
        self.max_depth = max_depth
        self.max_features = max_features
        self.min_samples_leaf = min_samples_leaf
        self.n_iter = n_iter
        self.metric = metric
        self.xgb_max_n_estimators = xgb_max_n_estimators
        self.xgb_n_estimators_step = xgb_n_estimators_step
        self.xgb_max_depth = xgb_max_depth
        self.xgb_min_learning_rate = xgb_min_learning_rate
        self.xgb_max_learning_rate = xgb_max_learning_rate
        self.xgb_learning_rate_step = xgb_learning_rate_step
        self.xgb_booster = xgb_booster
        self.xgb_min_gamma = xgb_min_gamma
        self.xgb_max_gamma = xgb_max_gamma
        self.xgb_gamma_step = xgb_gamma_step
        self.xgb_min_min_child_weight = xgb_min_min_child_weight
        self.xgb_max_min_child_weight = xgb_max_min_child_weight
        self.xgb_min_child_weight_step = xgb_min_child_weight_step
        self.xgb_subsample = xgb_subsample
        self.xgb_colsample_bytree = xgb_colsample_bytree
        self.xgb_colsample_bylevel = xgb_colsample_bylevel
        self.xgb_colsample_bynode = xgb_colsample_bynode
        self.xgb_min_reg_alpha = xgb_min_reg_alpha
        self.xgb_max_reg_alpha = xgb_max_reg_alpha
        self.xgb_reg_alpha_step = xgb_reg_alpha_step
        self.xgb_min_reg_lambda = xgb_min_reg_lambda
        self.xgb_max_reg_lambda = xgb_max_reg_lambda
        self.xgb_reg_lambda_step = xgb_reg_lambda_step
        self.n_jobs = n_jobs
        self.random_state = random_state
        self.verbose = verbose

    def fit(self, X, y):
        """
        Build a LCE tree from the training set (X, y).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The training input samples.

        y : array-like of shape (n_samples,)
            The target values (real numbers).

        Returns
        -------
        self : object
        """
        self.n_features_in_ = X.shape[1]

        def _build_tree(X, y):
            """Build a LCE tree."""
            global index_node_global

            def _create_node(X, y, depth, container):
                """Create a node in the tree."""
                # Add XGBoost predictions as features to the dataset
                model_node = xgb_opt_regressor(
                    X,
                    y,
                    n_iter=self.n_iter,
                    metric=self.metric,
                    n_estimators=self.xgb_max_n_estimators,
                    n_estimators_step=self.xgb_n_estimators_step,
                    max_depth=self.xgb_max_depth,
                    min_learning_rate=self.xgb_min_learning_rate,
                    max_learning_rate=self.xgb_max_learning_rate,
                    learning_rate_step=self.xgb_learning_rate_step,
                    booster=self.xgb_booster,
                    min_gamma=self.xgb_min_gamma,
                    max_gamma=self.xgb_max_gamma,
                    gamma_step=self.xgb_gamma_step,
                    min_min_child_weight=self.xgb_min_min_child_weight,
                    max_min_child_weight=self.xgb_max_min_child_weight,
                    min_child_weight_step=self.xgb_min_child_weight_step,
                    subsample=self.xgb_subsample,
                    colsample_bytree=self.xgb_colsample_bytree,
                    colsample_bylevel=self.xgb_colsample_bylevel,
                    colsample_bynode=self.xgb_colsample_bynode,
                    min_reg_alpha=self.xgb_min_reg_alpha,
                    max_reg_alpha=self.xgb_max_reg_alpha,
                    reg_alpha_step=self.xgb_reg_alpha_step,
                    min_reg_lambda=self.xgb_min_reg_lambda,
                    max_reg_lambda=self.xgb_max_reg_lambda,
                    reg_lambda_step=self.xgb_reg_lambda_step,
                    n_jobs=self.n_jobs,
                    random_state=self.random_state,
                )
                preds = np.around(model_node.predict(X), 6)
                X = np.insert(X, X.shape[1], 0, axis=1)
                X[:, -1] = preds

                # Missing data information
                num_nans = np.isnan(X).any(axis=1).sum()
                if num_nans > 0:
                    missing = True
                    if num_nans == y.size:
                        missing_only = True
                    else:
                        missing_only = False
                else:
                    missing = False
                    missing_only = False

                # Split
                split_val_conditions = [y.size > 1, missing_only == False]
                if all(split_val_conditions):
                    split = DecisionTreeRegressor(
                        criterion=self.criterion,
                        splitter=self.splitter,
                        max_depth=1,
                        max_features=self.max_features,
                        random_state=self.random_state,
                    )
                    if missing:
                        nans = np.isnan(X).any(axis=1)
                        split.fit(X[~nans], y[~nans])
                    else:
                        split.fit(X, y)
                else:
                    split = None

                # Node information
                node = {
                    "index": container["index_node_global"],
                    "model": model_node,
                    "data": (X, y),
                    "split": split,
                    "missing": {"missing": missing, "missing_only": missing_only},
                    "missing_side": None,
                    "children": {"left": None, "right": None},
                    "depth": depth,
                }
                container["index_node_global"] += 1
                return node

            def _splitter(node):
                """Perform the split of a node."""
                # Extract data
                X, y = node["data"]
                depth = node["depth"]
                split = node["split"]
                missing = node["missing"]["missing"]
                missing_only = node["missing"]["missing_only"]

                did_split = False
                data = None

                # Perform split if the conditions are met
                stopping_criteria = [
                    depth >= 0,
                    depth < self.max_depth,
                    X[:, 0].size > 1,
                    missing_only == False,
                ]

                if all(stopping_criteria):
                    if missing:
                        nans = np.isnan(X).any(axis=1)
                        X_withoutnans, y_withoutnans = X[~nans], y[~nans]
                        leafs = split.apply(X_withoutnans)
                        (X_left, y_left), (X_right, y_right) = (
                            np.squeeze(X_withoutnans[np.argwhere(leafs == 1), :]),
                            np.squeeze(y_withoutnans[np.argwhere(leafs == 1)]),
                        ), (
                            np.squeeze(X_withoutnans[np.argwhere(leafs == 2), :]),
                            np.squeeze(y_withoutnans[np.argwhere(leafs == 2)]),
                        )
                    else:
                        leafs = split.apply(X)
                        (X_left, y_left), (X_right, y_right) = (
                            np.squeeze(X[np.argwhere(leafs == 1), :]),
                            np.squeeze(y[np.argwhere(leafs == 1)]),
                        ), (
                            np.squeeze(X[np.argwhere(leafs == 2), :]),
                            np.squeeze(y[np.argwhere(leafs == 2)]),
                        )

                    N_left, N_right = y_left.size, y_right.size

                    split_conditions = [
                        N_left >= self.min_samples_leaf,
                        N_right >= self.min_samples_leaf,
                    ]

                    if all(split_conditions):
                        did_split = True

                        if N_left == 1:
                            X_left = X_left.reshape(-1, 1).T
                            node["missing_side"] = "left"
                            if missing:
                                X_left = np.append(X_left, X[nans], axis=0)
                                y_left = np.append([y_left], y[nans], axis=0)

                        if N_right == 1:
                            X_right = X_right.reshape(-1, 1).T
                            if N_left > 1:
                                node["missing_side"] = "right"
                                if missing:
                                    X_right = np.append(X_right, X[nans], axis=0)
                                    y_right = np.append([y_right], y[nans], axis=0)

                        score_conditions = [N_left > 1, N_right > 1]
                        if all(score_conditions):
                            if split.score(X_left, y_left) > split.score(
                                X_right, y_right
                            ):
                                node["missing_side"] = "left"
                                if missing:
                                    X_left = np.append(X_left, X[nans], axis=0)
                                    y_left = np.append(y_left, y[nans], axis=0)
                            else:
                                node["missing_side"] = "right"
                                if missing:
                                    X_right = np.append(X_right, X[nans], axis=0)
                                    y_right = np.append(y_right, y[nans], axis=0)

                        data = [(X_left, y_left), (X_right, y_right)]

                result = {"did_split": did_split, "data": data}
                return result

            def _split_traverse_node(node, container):
                """Process splitting results and continue with child nodes."""
                # Perform split and collect result
                result = _splitter(node)

                # Return terminal node if no split
                if not result["did_split"]:
                    if self.verbose > 0 and self.n_jobs == None:
                        depth_spacing_str = " ".join([" "] * node["depth"])
                        print(
                            " {}*leaf {} @ depth {}: Unique_y {},  N_samples {}".format(
                                depth_spacing_str,
                                node["index"],
                                node["depth"],
                                np.unique(node["data"][1]),
                                np.unique(node["data"][1], return_counts=True)[1],
                            )
                        )
                    return
                del node["data"]

                # Extract splitting results
                (X_left, y_left), (X_right, y_right) = result["data"]

                # Report created node to user
                if self.verbose > 0 and self.n_jobs == None:
                    depth_spacing_str = " ".join([" "] * node["depth"])
                    print(
                        " {}node {} @ depth {}: dataset={}, N_left={}, N_right={}".format(
                            depth_spacing_str,
                            node["index"],
                            node["depth"],
                            (X_left.shape[0] + X_right.shape[0], X_left.shape[1]),
                            X_left.shape[0],
                            X_right.shape[0],
                        )
                    )

                # Create child nodes
                node["children"]["left"] = _create_node(
                    X_left, y_left, node["depth"] + 1, container
                )
                node["children"]["right"] = _create_node(
                    X_right, y_right, node["depth"] + 1, container
                )

                # Split nodes
                _split_traverse_node(node["children"]["left"], container)
                _split_traverse_node(node["children"]["right"], container)

            container = {"index_node_global": 0}
            if self.verbose > 0 and self.n_jobs == None:
                print("\nNew Tree")
            root = _create_node(X, y, 0, container)
            _split_traverse_node(root, container)
            return root

        self.tree = _build_tree(X, y)
        return self

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

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The training input samples.

        Returns
        -------
        y : ndarray of shape (n_samples,)
            The predicted values.
        """

        def _base(node, X):
            y_pred = np.around(node["model"].predict(X[:, 1:]), 6)
            X = np.insert(X, X.shape[1], 0, axis=1)
            X[:, -1] = y_pred
            return X

        def _predict(node, X, y_pred_final=None):
            if X.ndim == 1:
                X = _base(node, X.reshape(-1, 1).T)
            else:
                X = _base(node, X)

            no_children = (
                node["children"]["left"] is None and node["children"]["right"] is None
            )
            if no_children:
                y_pred = np.column_stack((X[:, :1], X[:, -1:]))
                if y_pred_final is not None:
                    y_pred_final = np.concatenate((y_pred_final, y_pred), axis=0)
                else:
                    y_pred_final = y_pred
                return y_pred_final

            else:
                if np.isnan(X).sum() > 0:
                    nans = np.isnan(X).any(axis=1)
                    leafs = node["split"].apply(X[~nans, 1:])
                    X_left, X_right = np.squeeze(
                        X[~nans][np.argwhere(leafs == 1), :]
                    ), np.squeeze(X[~nans][np.argwhere(leafs == 2), :])
                    if node["missing_side"] == "left":
                        X_left, X_right = (
                            np.concatenate((X_left, X[nans]), axis=0),
                            X_right,
                        )
                    else:
                        X_left, X_right = X_left, np.concatenate(
                            (X_right, X[nans]), axis=0
                        )
                else:
                    leafs = node["split"].apply(X[:, 1:])
                    X_left, X_right = np.squeeze(
                        X[np.argwhere(leafs == 1), :]
                    ), np.squeeze(X[np.argwhere(leafs == 2), :])

                if len(X_left) > 0:
                    y_pred_final = _predict(
                        node["children"]["left"], X_left, y_pred_final
                    )
                if len(X_right) > 0:
                    y_pred_final = _predict(
                        node["children"]["right"], X_right, y_pred_final
                    )
                return y_pred_final

        index = np.arange(0, X.shape[0]).reshape(-1, 1)
        X = np.concatenate((index, X), axis=1)
        y_pred = _predict(self.tree, X, None)
        y_pred = y_pred[y_pred[:, 0].argsort()]
        y_pred = np.squeeze(y_pred[:, 1:])
        return y_pred

    def set_params(self, **params):
        """
        Set the parameters of the estimator.

        Parameters
        ----------
        **params : dict
            Estimator parameters.

        Returns
        -------
        self : object
        """
        if not params:
            return self

        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                self.kwargs[key] = value
        return self
