from abc import ABCMeta, abstractmethod
from functools import partial
import threading
import warnings

from joblib import Parallel, delayed
import numpy as np
from sklearn.ensemble._base import _partition_estimators
from sklearn.ensemble._forest import (
    BaseForest,
    _accumulate_prediction,
    _generate_unsampled_indices,
    _get_n_samples_bootstrap,
    _parallel_build_trees,
)
from sklearn.tree._tree import DTYPE
from sklearn.utils._tags import get_tags
from sklearn.utils.validation import check_is_fitted, check_random_state, validate_data

from ..base import SurvivalAnalysisMixin
from ..docstrings import append_cumulative_hazard_example, append_survival_function_example
from ..metrics import concordance_index_censored
from ..tree import ExtraSurvivalTree, SurvivalTree
from ..tree._criterion import get_unique_times
from ..tree.tree import _array_to_step_function
from ..util import check_array_survival

__all__ = ["RandomSurvivalForest", "ExtraSurvivalTrees"]

MAX_INT = np.iinfo(np.int32).max


def _sklearn_tags_patch(self):
    # BaseForest.__sklearn_tags__ calls
    # type(self.estimator)(criterion=self.criterions),
    # which is incompatible with LogrankCriterion
    if isinstance(self, _BaseSurvivalForest):
        estimator = type(self.estimator)()
    else:
        estimator = type(self.estimator)(criterion=self.criterion)
    tags = super(BaseForest, self).__sklearn_tags__()
    tags.input_tags.allow_nan = get_tags(estimator).input_tags.allow_nan
    return tags


BaseForest.__sklearn_tags__ = _sklearn_tags_patch


class _BaseSurvivalForest(BaseForest, metaclass=ABCMeta):
    """
    Base class for forest-based estimators for survival analysis.

    Warning: This class should not be used directly. Use derived classes
    instead.
    """

    @abstractmethod
    def __init__(
        self,
        estimator,
        n_estimators=100,
        *,
        estimator_params=tuple(),
        bootstrap=False,
        oob_score=False,
        n_jobs=None,
        random_state=None,
        verbose=0,
        warm_start=False,
        max_samples=None,
    ):
        super().__init__(
            estimator,
            n_estimators=n_estimators,
            estimator_params=estimator_params,
            bootstrap=bootstrap,
            oob_score=oob_score,
            n_jobs=n_jobs,
            random_state=random_state,
            verbose=verbose,
            warm_start=warm_start,
            class_weight=None,
            max_samples=max_samples,
        )

    @property
    def feature_importances_(self):
        """Not implemented"""
        raise NotImplementedError()

    def fit(self, X, y, sample_weight=None):
        """Build a forest of survival trees from the training set (X, y).

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix

        y : structured array, shape = (n_samples,)
            A structured array with two fields. The first field is a boolean
            where ``True`` indicates an event and ``False`` indicates right-censoring.
            The second field is a float with the time of event or time of censoring.

        Returns
        -------
        self
        """
        self._validate_params()

        X = validate_data(self, X, dtype=DTYPE, accept_sparse="csc", ensure_min_samples=2, ensure_all_finite=False)
        event, time = check_array_survival(X, y)

        # _compute_missing_values_in_feature_mask checks if X has missing values and
        # will raise an error if the underlying tree base estimator can't handle missing
        # values.
        estimator = type(self.estimator)()
        missing_values_in_feature_mask = estimator._compute_missing_values_in_feature_mask(
            X, estimator_name=self.__class__.__name__
        )

        self._n_samples, self.n_features_in_ = X.shape
        time = time.astype(np.float64)
        self.unique_times_, self.is_event_time_ = get_unique_times(time, event)
        self.n_outputs_ = self.unique_times_.shape[0]

        y_numeric = np.empty((X.shape[0], 2), dtype=np.float64)
        y_numeric[:, 0] = time
        y_numeric[:, 1] = event.astype(np.float64)

        # Get bootstrap sample size
        if not self.bootstrap and self.max_samples is not None:  # pylint: disable=no-else-raise
            raise ValueError(
                "`max_sample` cannot be set if `bootstrap=False`. "
                "Either switch to `bootstrap=True` or set "
                "`max_sample=None`."
            )
        elif self.bootstrap:
            n_samples_bootstrap = _get_n_samples_bootstrap(n_samples=X.shape[0], max_samples=self.max_samples)
        else:
            n_samples_bootstrap = None

        self._n_samples_bootstrap = n_samples_bootstrap

        # Check parameters
        self._validate_estimator()

        if not self.bootstrap and self.oob_score:
            raise ValueError("Out of bag estimation only available if bootstrap=True")

        random_state = check_random_state(self.random_state)

        if not self.warm_start or not hasattr(self, "estimators_"):
            # Free allocated memory, if any
            self.estimators_ = []

        n_more_estimators = self.n_estimators - len(self.estimators_)

        if n_more_estimators < 0:  # pylint: disable=no-else-raise
            raise ValueError(
                f"n_estimators={self.n_estimators} must be larger or equal to "
                f"len(estimators_)={len(self.estimators_)} when warm_start==True"
            )
        elif n_more_estimators == 0:
            warnings.warn("Warm-start fitting without increasing n_estimators does not fit new trees.", stacklevel=2)
        else:
            if self.warm_start and len(self.estimators_) > 0:
                # We draw from the random state to get the random state we
                # would have got if we hadn't used a warm_start.
                random_state.randint(MAX_INT, size=len(self.estimators_))

            trees = [self._make_estimator(append=False, random_state=random_state) for i in range(n_more_estimators)]

            y_tree = (
                y_numeric,
                self.unique_times_,
                self.is_event_time_,
            )
            # Parallel loop: we prefer the threading backend as the Cython code
            # for fitting the trees is internally releasing the Python GIL
            # making threading more efficient than multiprocessing in
            # that case. However, for joblib 0.12+ we respect any
            # parallel_backend contexts set at a higher level,
            # since correctness does not rely on using threads.
            trees = Parallel(n_jobs=self.n_jobs, verbose=self.verbose, prefer="threads")(
                delayed(_parallel_build_trees)(
                    t,
                    self.bootstrap,
                    X,
                    y_tree,
                    sample_weight,
                    i,
                    len(trees),
                    verbose=self.verbose,
                    n_samples_bootstrap=n_samples_bootstrap,
                    missing_values_in_feature_mask=missing_values_in_feature_mask,
                )
                for i, t in enumerate(trees)
            )

            # Collect newly grown trees
            self.estimators_.extend(trees)

        if self.oob_score:
            self._set_oob_score_and_attributes(X, (event, time))

        return self

    def _set_oob_score_and_attributes(self, X, y):
        """Calculate out of bag predictions and score."""
        n_samples = X.shape[0]
        event, time = y

        predictions = np.zeros(n_samples)
        n_predictions = np.zeros(n_samples)

        n_samples_bootstrap = _get_n_samples_bootstrap(n_samples, self.max_samples)

        for estimator in self.estimators_:
            unsampled_indices = _generate_unsampled_indices(estimator.random_state, n_samples, n_samples_bootstrap)
            p_estimator = estimator.predict(X[unsampled_indices, :], check_input=False)

            predictions[unsampled_indices] += p_estimator
            n_predictions[unsampled_indices] += 1

        if (n_predictions == 0).any():
            warnings.warn(
                "Some inputs do not have OOB scores. "
                "This probably means too few trees were used "
                "to compute any reliable oob estimates.",
                stacklevel=3,
            )
            n_predictions[n_predictions == 0] = 1

        predictions /= n_predictions
        self.oob_prediction_ = predictions

        self.oob_score_ = concordance_index_censored(event, time, predictions)[0]

    def _predict(self, predict_fn, X):
        check_is_fitted(self, "estimators_")
        # Check data
        X = self._validate_X_predict(X)

        # Assign chunk of trees to jobs
        n_jobs, _, _ = _partition_estimators(self.n_estimators, self.n_jobs)

        # avoid storing the output of every estimator by summing them here
        if predict_fn == "predict":
            y_hat = np.zeros((X.shape[0]), dtype=np.float64)
        else:
            y_hat = np.zeros((X.shape[0], self.n_outputs_), dtype=np.float64)

        def _get_fn(est, name):
            fn = getattr(est, name)
            if name in ("predict_cumulative_hazard_function", "predict_survival_function"):
                fn = partial(fn, return_array=True)
            return fn

        # Parallel loop
        lock = threading.Lock()
        Parallel(n_jobs=n_jobs, verbose=self.verbose, require="sharedmem")(
            delayed(_accumulate_prediction)(_get_fn(e, predict_fn), X, [y_hat], lock) for e in self.estimators_
        )

        y_hat /= len(self.estimators_)

        return y_hat

    def predict(self, X):
        r"""Predict risk score.

        The ensemble risk score is the total number of events,
        which can be estimated by the sum of the estimated
        ensemble cumulative hazard function :math:`\hat{H}_e`.

        .. math::

            \sum_{j=1}^{n} \hat{H}_e(T_{j} \mid x) ,

        where :math:`n` denotes the total number of distinct
        event times in the training data.

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix.

        Returns
        -------
        risk_scores : ndarray, shape = (n_samples,)
            Predicted risk scores.
        """
        return self._predict("predict", X)

    def predict_cumulative_hazard_function(self, X, return_array=False):
        arr = self._predict("predict_cumulative_hazard_function", X)
        if return_array:
            return arr
        return _array_to_step_function(self.unique_times_, arr)

    def predict_survival_function(self, X, return_array=False):
        arr = self._predict("predict_survival_function", X)
        if return_array:
            return arr
        return _array_to_step_function(self.unique_times_, arr)


class RandomSurvivalForest(SurvivalAnalysisMixin, _BaseSurvivalForest):
    """A random survival forest.

    A random survival forest is a meta estimator that fits a number of
    survival trees on various sub-samples of the dataset and uses
    averaging to improve the predictive accuracy and control over-fitting.
    The sub-sample size is always the same as the original input sample
    size but the samples are drawn with replacement if
    `bootstrap=True` (default).

    In each survival tree, the quality of a split is measured by the
    log-rank splitting rule.

    See the :ref:`User Guide </user_guide/random-survival-forest.ipynb>`,
    [1]_ and [2]_ for further description.

    Parameters
    ----------
    n_estimators : int, optional, default: 100
        The number of trees in the forest.

    max_depth : int or None, optional, default: None
        The maximum depth of the tree. If None, then nodes are expanded until
        all leaves are pure or until all leaves contain less than
        min_samples_split samples.

    min_samples_split : int, float, optional, default: 6
        The minimum number of samples required to split an internal node:

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

    min_samples_leaf : int, float, optional, default: 3
        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.  This may have the effect of smoothing the model,
        especially in regression.

        - 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.

    min_weight_fraction_leaf : float, optional, default: 0.
        The minimum weighted fraction of the sum total of weights (of all
        the input samples) required to be at a leaf node. Samples have
        equal weight when sample_weight is not provided.

    max_features : int, float, {'sqrt', 'log2'} or None, optional, default: 'sqrt'
        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
          `int(max_features * n_features)` features are considered at each
          split.
        - If "sqrt", then `max_features=sqrt(n_features)`.
        - 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.

    max_leaf_nodes : int or None, optional, default: None
        Grow a tree with ``max_leaf_nodes`` in best-first fashion.
        Best nodes are defined as relative reduction in impurity.
        If None then unlimited number of leaf nodes.

    bootstrap : bool, optional, default: True
        Whether bootstrap samples are used when building trees. If False, the
        whole dataset is used to build each tree.

    oob_score : bool, optional, default: False
        Whether to use out-of-bag samples to estimate
        the generalization accuracy.

    n_jobs : int or None, optional, default: None
        The number of jobs to run in parallel. :meth:`fit`, :meth:`predict`,
        :meth:`decision_path` and :meth:`apply` are all parallelized over the
        trees. ``None`` means 1 unless in a :obj:`joblib.parallel_backend`
        context. ``-1`` means using all processors.

    random_state : int, RandomState instance or None, optional, default: None
        Controls both the randomness of the bootstrapping of the samples used
        when building trees (if ``bootstrap=True``) and the sampling of the
        features to consider when looking for the best split at each node
        (if ``max_features < n_features``).

    verbose : int, optional, default: 0
        Controls the verbosity when fitting and predicting.

    warm_start : bool, optional, default: False
        When set to ``True``, reuse the solution of the previous call to fit
        and add more estimators to the ensemble, otherwise, just fit a whole
        new forest.

    max_samples : int or float, optional, default: None
        If bootstrap is True, the number of samples to draw from X
        to train each base estimator.

        - If None (default), then draw `X.shape[0]` samples.
        - If int, then draw `max_samples` samples.
        - If float, then draw `max_samples * X.shape[0]` samples. Thus,
          `max_samples` should be in the interval `(0.0, 1.0]`.

    low_memory : bool, optional, default: False
        If set, :meth:`predict` computations use reduced memory but :meth:`predict_cumulative_hazard_function`
        and :meth:`predict_survival_function` are not implemented.

    Attributes
    ----------
    estimators_ : list of SurvivalTree instances
        The collection of fitted sub-estimators.

    unique_times_ : ndarray, shape = (n_unique_times,)
        Unique time points.

    n_features_in_ : int
        Number of features seen during ``fit``.

    feature_names_in_ : ndarray, shape = (`n_features_in_`,)
        Names of features seen during ``fit``. Defined only when `X`
        has feature names that are all strings.

    oob_score_ : float
        Concordance index of the training dataset obtained
        using an out-of-bag estimate.

    See also
    --------
    sksurv.tree.SurvivalTree
        A single survival tree.

    Notes
    -----
    The default values for the parameters controlling the size of the trees
    (e.g. ``max_depth``, ``min_samples_leaf``, etc.) lead to fully grown and
    unpruned trees which can potentially be very large on some data sets. To
    reduce memory consumption, the complexity and size of the trees should be
    controlled by setting those parameter values.

    Compared to scikit-learn's random forest models, :class:`RandomSurvivalForest`
    currently does not support controlling the depth of a tree based on the log-rank
    test statistics or it's associated p-value, i.e., the parameters
    `min_impurity_decrease` or `min_impurity_split` are absent.
    In addition, the `feature_importances_` attribute is not available.
    It is recommended to estimate feature importances via
    :func:`sklearn.inspection.permutation_importance`.

    The features are always randomly permuted at each split. Therefore,
    the best found split may vary, even with the same training data,
    ``max_features=n_features`` and ``bootstrap=False``, if the improvement
    of the criterion is identical for several splits enumerated during the
    search of the best split. To obtain a deterministic behavior during
    fitting, ``random_state`` has to be fixed.

    References
    ----------
    .. [1] Ishwaran, H., Kogalur, U. B., Blackstone, E. H., & Lauer, M. S. (2008).
           Random survival forests. The Annals of Applied Statistics, 2(3), 841–860.

    .. [2] Ishwaran, H., Kogalur, U. B. (2007). Random survival forests for R.
           R News, 7(2), 25–31. https://cran.r-project.org/doc/Rnews/Rnews_2007-2.pdf.
    """

    _parameter_constraints = {
        **BaseForest._parameter_constraints,
        **SurvivalTree._parameter_constraints,
    }
    _parameter_constraints.pop("splitter")

    def __init__(
        self,
        n_estimators=100,
        *,
        max_depth=None,
        min_samples_split=6,
        min_samples_leaf=3,
        min_weight_fraction_leaf=0.0,
        max_features="sqrt",
        max_leaf_nodes=None,
        bootstrap=True,
        oob_score=False,
        n_jobs=None,
        random_state=None,
        verbose=0,
        warm_start=False,
        max_samples=None,
        low_memory=False,
    ):
        super().__init__(
            estimator=SurvivalTree(),
            n_estimators=n_estimators,
            estimator_params=(
                "max_depth",
                "min_samples_split",
                "min_samples_leaf",
                "min_weight_fraction_leaf",
                "max_features",
                "max_leaf_nodes",
                "random_state",
                "low_memory",
            ),
            bootstrap=bootstrap,
            oob_score=oob_score,
            n_jobs=n_jobs,
            random_state=random_state,
            verbose=verbose,
            warm_start=warm_start,
            max_samples=max_samples,
        )

        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.low_memory = low_memory

    @append_cumulative_hazard_example(estimator_mod="ensemble", estimator_class="RandomSurvivalForest")
    def predict_cumulative_hazard_function(self, X, return_array=False):
        """Predict cumulative hazard function.

        For each tree in the ensemble, the cumulative hazard
        function (CHF) for an individual with feature vector
        :math:`x` is computed from all samples of the bootstrap
        sample that are in the same terminal node as :math:`x`.
        It is estimated by the Nelson–Aalen estimator.
        The ensemble CHF at time :math:`t` is the average
        value across all trees in the ensemble at the
        specified time point.

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix.

        return_array : bool, default: False
            Whether to return a single array of cumulative hazard values
            or a list of step functions.

            If `False`, a list of :class:`sksurv.functions.StepFunction`
            objects is returned.

            If `True`, a 2d-array of shape `(n_samples, n_unique_times)` is
            returned, where `n_unique_times` is the number of unique
            event times in the training data. Each row represents the cumulative
            hazard function of an individual evaluated at `unique_times_`.

        Returns
        -------
        cum_hazard : ndarray
            If `return_array` is `False`, an array of `n_samples`
            :class:`sksurv.functions.StepFunction` instances is returned.

            If `return_array` is `True`, a numeric array of shape
            `(n_samples, n_unique_times_)` is returned.

        Examples
        --------
        """
        return super().predict_cumulative_hazard_function(X, return_array)

    @append_survival_function_example(estimator_mod="ensemble", estimator_class="RandomSurvivalForest")
    def predict_survival_function(self, X, return_array=False):
        """Predict survival function.

        For each tree in the ensemble, the survival function
        for an individual with feature vector :math:`x` is
        computed from all samples of the bootstrap sample that
        are in the same terminal node as :math:`x`.
        It is estimated by the Kaplan-Meier estimator.
        The ensemble survival function at time :math:`t` is
        the average value across all trees in the ensemble at
        the specified time point.

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix.

        return_array : bool, default: False
            Whether to return a single array of survival probabilities
            or a list of step functions.

            If `False`, a list of :class:`sksurv.functions.StepFunction`
            objects is returned.

            If `True`, a 2d-array of shape `(n_samples, n_unique_times)` is
            returned, where `n_unique_times` is the number of unique
            event times in the training data. Each row represents the survival
            function of an individual evaluated at `unique_times_`.

        Returns
        -------
        survival : ndarray
            If `return_array` is `False`, an array of `n_samples`
            :class:`sksurv.functions.StepFunction` instances is returned.

            If `return_array` is `True`, a numeric array of shape
            `(n_samples, n_unique_times_)` is returned.

        Examples
        --------
        """
        return super().predict_survival_function(X, return_array)


class ExtraSurvivalTrees(SurvivalAnalysisMixin, _BaseSurvivalForest):
    """An extremely random survival forest.

    This class implements a meta estimator that fits a number of randomized
    survival trees (a.k.a. extra-trees) on various sub-samples of the dataset
    and uses averaging to improve the predictive accuracy and control
    over-fitting. The sub-sample size is always the same as the original
    input sample size but the samples are drawn with replacement if
    `bootstrap=True` (default).

    In each randomized survival tree, the quality of a split is measured by
    the log-rank splitting rule.

    Compared to :class:`RandomSurvivalForest`, randomness goes one step
    further in the way splits are computed. As in
    :class:`RandomSurvivalForest`, a random subset of candidate features is
    used, but instead of looking for the most discriminative thresholds,
    thresholds are drawn at random for each candidate feature and the best of
    these randomly-generated thresholds is picked as the splitting rule.

    Parameters
    ----------
    n_estimators : int, optional, default: 100
        The number of trees in the forest.

    max_depth : int or None, optional, default: None
        The maximum depth of the tree. If None, then nodes are expanded until
        all leaves are pure or until all leaves contain less than
        min_samples_split samples.

    min_samples_split : int, float, optional, default: 6
        The minimum number of samples required to split an internal node:

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

    min_samples_leaf : int, float, optional, default: 3
        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.  This may have the effect of smoothing the model,
        especially in regression.

        - 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.

    min_weight_fraction_leaf : float, optional, default: 0.
        The minimum weighted fraction of the sum total of weights (of all
        the input samples) required to be at a leaf node. Samples have
        equal weight when sample_weight is not provided.

    max_features : int, float, {'sqrt', 'log2'} or None, optional, default: 'sqrt'
        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
          `int(max_features * n_features)` features are considered at each
          split.
        - If "sqrt", then `max_features=sqrt(n_features)`.
        - 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.

    max_leaf_nodes : int or None, optional, default: None
        Grow a tree with ``max_leaf_nodes`` in best-first fashion.
        Best nodes are defined as relative reduction in impurity.
        If None then unlimited number of leaf nodes.

    bootstrap : bool, optional, default: True
        Whether bootstrap samples are used when building trees. If False, the
        whole dataset is used to build each tree.

    oob_score : bool, optional, default: False
        Whether to use out-of-bag samples to estimate
        the generalization accuracy.

    n_jobs : int or None, optional, default: None
        The number of jobs to run in parallel. :meth:`fit`, :meth:`predict`,
        :meth:`decision_path` and :meth:`apply` are all parallelized over the
        trees. ``None`` means 1 unless in a :obj:`joblib.parallel_backend`
        context. ``-1`` means using all processors.

    random_state : int, RandomState instance or None, optional, default: None
        Controls both the randomness of the bootstrapping of the samples used
        when building trees (if ``bootstrap=True``) and the sampling of the
        features to consider when looking for the best split at each node
        (if ``max_features < n_features``).

    verbose : int, optional, default: 0
        Controls the verbosity when fitting and predicting.

    warm_start : bool, optional, default: False
        When set to ``True``, reuse the solution of the previous call to fit
        and add more estimators to the ensemble, otherwise, just fit a whole
        new forest.

    max_samples : int or float, optional, default: None
        If bootstrap is True, the number of samples to draw from X
        to train each base estimator.

        - If None (default), then draw `X.shape[0]` samples.
        - If int, then draw `max_samples` samples.
        - If float, then draw `max_samples * X.shape[0]` samples. Thus,
          `max_samples` should be in the interval `(0.0, 1.0]`.

    low_memory : bool, optional, default: False
        If set, :meth:`predict` computations use reduced memory but :meth:`predict_cumulative_hazard_function`
        and :meth:`predict_survival_function` are not implemented.

    Attributes
    ----------
    estimators_ : list of SurvivalTree instances
        The collection of fitted sub-estimators.

    unique_times_ : ndarray, shape = (n_unique_times,)
        Unique time points.

    n_features_in_ : int
        Number of features seen during ``fit``.

    feature_names_in_ : ndarray, shape = (`n_features_in_`,)
        Names of features seen during ``fit``. Defined only when `X`
        has feature names that are all strings.

    oob_score_ : float
        Concordance index of the training dataset obtained
        using an out-of-bag estimate.

    See also
    --------
    sksurv.tree.SurvivalTree
        A single survival tree.
    """

    _parameter_constraints = {
        **BaseForest._parameter_constraints,
        **SurvivalTree._parameter_constraints,
    }
    _parameter_constraints.pop("splitter")

    def __init__(
        self,
        n_estimators=100,
        *,
        max_depth=None,
        min_samples_split=6,
        min_samples_leaf=3,
        min_weight_fraction_leaf=0.0,
        max_features="sqrt",
        max_leaf_nodes=None,
        bootstrap=True,
        oob_score=False,
        n_jobs=None,
        random_state=None,
        verbose=0,
        warm_start=False,
        max_samples=None,
        low_memory=False,
    ):
        super().__init__(
            estimator=ExtraSurvivalTree(),
            n_estimators=n_estimators,
            estimator_params=(
                "max_depth",
                "min_samples_split",
                "min_samples_leaf",
                "min_weight_fraction_leaf",
                "max_features",
                "max_leaf_nodes",
                "random_state",
                "low_memory",
            ),
            bootstrap=bootstrap,
            oob_score=oob_score,
            n_jobs=n_jobs,
            random_state=random_state,
            verbose=verbose,
            warm_start=warm_start,
            max_samples=max_samples,
        )

        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.low_memory = low_memory

    @append_cumulative_hazard_example(estimator_mod="ensemble", estimator_class="ExtraSurvivalTrees")
    def predict_cumulative_hazard_function(self, X, return_array=False):
        """Predict cumulative hazard function.

        For each tree in the ensemble, the cumulative hazard
        function (CHF) for an individual with feature vector
        :math:`x` is computed from all samples of the bootstrap
        sample that are in the same terminal node as :math:`x`.
        It is estimated by the Nelson–Aalen estimator.
        The ensemble CHF at time :math:`t` is the average
        value across all trees in the ensemble at the
        specified time point.

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix.

        return_array : bool, default: False
            Whether to return a single array of cumulative hazard values
            or a list of step functions.

            If `False`, a list of :class:`sksurv.functions.StepFunction`
            objects is returned.

            If `True`, a 2d-array of shape `(n_samples, n_unique_times)` is
            returned, where `n_unique_times` is the number of unique
            event times in the training data. Each row represents the cumulative
            hazard function of an individual evaluated at `unique_times_`.

        Returns
        -------
        cum_hazard : ndarray
            If `return_array` is `False`, an array of `n_samples`
            :class:`sksurv.functions.StepFunction` instances is returned.

            If `return_array` is `True`, a numeric array of shape
            `(n_samples, n_unique_times_)` is returned.

        Examples
        --------
        """
        return super().predict_cumulative_hazard_function(X, return_array)

    @append_survival_function_example(estimator_mod="ensemble", estimator_class="ExtraSurvivalTrees")
    def predict_survival_function(self, X, return_array=False):
        """Predict survival function.

        For each tree in the ensemble, the survival function
        for an individual with feature vector :math:`x` is
        computed from all samples of the bootstrap sample that
        are in the same terminal node as :math:`x`.
        It is estimated by the Kaplan-Meier estimator.
        The ensemble survival function at time :math:`t` is
        the average value across all trees in the ensemble at
        the specified time point.

        Parameters
        ----------
        X : array-like, shape = (n_samples, n_features)
            Data matrix.

        return_array : bool, default: False
            Whether to return a single array of survival probabilities
            or a list of step functions.

            If `False`, a list of :class:`sksurv.functions.StepFunction`
            objects is returned.

            If `True`, a 2d-array of shape `(n_samples, n_unique_times)` is
            returned, where `n_unique_times` is the number of unique
            event times in the training data. Each row represents the survival
            function of an individual evaluated at `unique_times_`.

        Returns
        -------
        survival : ndarray
            If `return_array` is `False`, an array of `n_samples`
            :class:`sksurv.functions.StepFunction` instances is returned.

            If `return_array` is `True`, a numeric array of shape
            `(n_samples, n_unique_times_)` is returned.

        Examples
        --------
        """
        return super().predict_survival_function(X, return_array)
