from itertools import combinations

import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from lifelines.statistics import logrank_test
from sklearn.base import clone, BaseEstimator
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, cross_val_score


class LateFusionClassifier(BaseEstimator):
    """
    Late fusion classifier for multimodal integration.

    Parameters
    ----------
    estimators: list of (str, estimator, list, dict) tuples.
        List of unimodal estimators to fit and fuse. Each unimodal estimator is associate with a tuple
        (`name`, `estimator` , `features`, `tune_dict`) where `name` is a string and corresponds to the name of the
        estimator, `estimator` ia scikit-learn estimator inheriting from BaseEstimator (can be a Pipeline), `features`
        is a list of indexes corresponding to the columns of the data associated with the modality of intereset, and
        `tune_dict` is either a dictionnary or a tuple (dict, n_iterations) for hyperparameter tuning and GridSearch or
        RandomSearch strategy respectively.

    cv: cross-validation generator
        cross-validation scheme for hyperparameter tuning (if `tuning` is not None) and/or calibration (if `calibration`
        is not None). The default is None

    score: str or callable.
        Score to use for tuning the unimodal models or weighting them at the late fusion step (i.e. sum of the unimodal
        predictions weighted by the performance of each unimodal model estimated with cross-validation).
        See sklearn.model_selection.cross_val_score for more details. The default is None.

    random_score: float.
        Random score for classification. Used when weighting the unimodal models with their estimated score. Weights
        will be max(score - random_score, 0). Unimodal models whose estimated performance is below the random_score
        will not be taken into account. The default is 0.5

    sup_weights: bool.
        Whether to use weights associated with the cross-validation performance of each unimodal model. If false no
        weights are used when fusing the unimodal predictions. The default is False.

    missing_threshold: float in ]0, 1].
        Minimum frequency of missing values to consider a whole modality missing (e.g., if `missing_threshold = 0.9` it
        means that for each sample and each modality at least 90% of the features associated with this modality must be
        missing to consider the whole modality missing). The default is 0.9.

    tuning: str or None.
        Strategy for tuning each model. Either 'gridsearch' for GridSearchCV or 'randomsearch' for RandomSearchCV. If
        None no hyperparameter tuning will be performed. The default is None.

    n_jobs: int.
        Number of jobs to run in parallel for hyperparameter tuning, collecting the predictions for calibration, or
        estimating the performance of each unimodal model with cross-validation. The default is None.

    calibration: str or None.
        Calibration strategy.
            * `calibration = 'late'` means that the fusion is made before calibration. The predictions of each
            multimodal combination are collected with cross-validation and a univariate logistic regression model is
            fitted to these predictions.
            * `calibration = 'early'` means that each unimodal model is calibrated prior to the late fusion. The
            unimodal predictions are collected with a cross-validation scheme and univariate logistic regression models
            are fitted.
            * `calibration = None` means that no calibration is performed.

    Attributes
    ----------
    best_params_: list of dict or empty list.
        List of best parameters for each unimodal predictor (output of GridSearchCV or RandomSearchCV). It follows the
        same order as the one of `estimators` list. If `tuning` is None returns an empty list (i.e., no hyperparameter
        tuning is performed).

    weights_: list of float.
        List of the weights associated to each modality and used at the late fusion stage for weighted sum.

    fitted_estimators_: list of estimators.
        List of fitted unimodal estimators.

    fitted_meta_estimators_: dictionary of estimators.
        Dictionary of meta-estimators for calibration. If `calibration = "early"` the keys correspond to the indexes of
        each unimodal estimator (i.e., from 0 to n_estimators-1) and the values correspond to the logistic regression
        estimators fitted to calibrate the unimodal models. If `calibration = "late"` the keys correspond to tuples
        characterizing each multimodal combination (e.g., (1, 3, 5)) and the values correspond th the logistic regression
        estimators fitted to clibrate the multimodal models.
    """

    def __init__(
        self,
        estimators,
        cv=None,
        score=None,
        random_score=0.5,
        sup_weights=False,
        missing_threshold=0.9,
        tuning=None,
        n_jobs=None,
        calibration="late",
    ):
        self.estimators = estimators
        self.cv = cv
        self.score = score
        self.random_score = random_score
        self.sup_weights = sup_weights
        self.missing_threshold = missing_threshold
        self.tuning = tuning
        self.n_jobs = n_jobs
        self.calibration = calibration

        self.best_params_ = []
        self.weights_ = []
        self.fitted_estimators_ = []
        self.fitted_meta_estimators_ = {}

    def fit(self, X, y):
        """
        Fit the latefusion classifier.

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        y: array of shape (n_samples,)
            Target to predict.

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        predictions = np.zeros((X.shape[0], len(self.estimators)))
        weights = np.zeros((X.shape[0], len(self.estimators)))
        i = 0
        for name, estim, features, grid in self.estimators:
            Xnew = X[:, features]
            # bool_mask = ~(np.sum(np.isnan(Xnew), axis=1) > self.missing_threshold * len(features))
            bool_mask = ~(
                np.sum(pd.isnull(Xnew), axis=1) > self.missing_threshold * len(features)
            )
            Xnew, ynew = Xnew[bool_mask, :], y[bool_mask]
            # Fit unimodal estimator
            self._fit_estim(
                Xnew, ynew, estim=estim, features=features, grid=grid, name=name
            )

            # Collect predictions and weights for further calibration
            if self.calibration is not None:
                weights[bool_mask, i] = max(self.weights_[-1] - self.random_score, 0)
                parallel = Parallel(n_jobs=self.n_jobs)
                collected_predictions = parallel(
                    delayed(_collect)(
                        Xdata=X,
                        ydata=y,
                        estimator=estim,
                        bmask=bool_mask,
                        feat=features,
                        train=train,
                        test=test,
                    )
                    for train, test in self.cv.split(X, y)
                )
                for indexes, preds in collected_predictions:
                    predictions[indexes, i] = preds
            i += 1

        # Calibrate models
        if self.calibration is not None:
            if self.calibration == "early":
                self._fit_early_calibration(predictions=predictions, y=y)
            elif self.calibration == "late":
                self._fit_late_calibration(
                    predictions=predictions, weights=weights, y=y
                )
            else:
                raise ValueError(
                    "'early', 'late' or None are the only values available for calibration parameter"
                )

        self.weights_ = np.array(self.weights_) - self.random_score
        self.weights_ = np.where(self.weights_ > 0, self.weights_, 0)
        return self

    def _fit_estim(self, X, y, estim, features, grid, name):
        """
        Fit a unimodal estimator.
        """
        if (self.tuning is not None) and (len(grid) > 0):
            if self.tuning == "gridsearch":
                search = GridSearchCV(
                    estimator=clone(estim),
                    param_grid=grid,
                    cv=self.cv,
                    scoring=self.score,
                    n_jobs=self.n_jobs,
                )
            elif self.tuning == "randomsearch":
                search = RandomizedSearchCV(
                    estimator=clone(estim),
                    param_distributions=grid[1],
                    n_iter=grid[0],
                    scoring=self.score,
                    n_jobs=self.n_jobs,
                    cv=self.cv,
                )

            search.fit(X, y)

            if self.sup_weights:
                self.weights_.append(search.best_score_)
            else:
                self.weights_.append(1.0)

            temp = search.best_estimator_
            self.best_params_.append(search.best_params_)
            # print("Best params " + name + " :", search.best_params_)
            # print("Best score " + name + " :", search.best_score_)
        else:
            if self.sup_weights:
                self.weights_.append(
                    np.mean(
                        cross_val_score(
                            estimator=clone(estim),
                            X=X,
                            y=y,
                            cv=self.cv,
                            scoring=self.score,
                            n_jobs=self.n_jobs,
                        )
                    )
                )
            else:
                self.weights_.append(1.0)
            temp = clone(estim).fit(X, y)

        self.fitted_estimators_.append((name, temp, features))
        return

    def _fit_early_calibration(self, predictions, y):
        """
        Calibrate only each unimodal predictor.
        """
        for i in range(len(self.estimators)):
            probas = predictions[:, i]
            mask = (probas > 0).reshape(-1)
            self.fitted_meta_estimators_[i] = LogisticRegression(
                class_weight="balanced"
            ).fit(probas[mask].reshape(-1, 1), y[mask])
        return

    def _fit_late_calibration(self, predictions, weights, y):
        """
        Calibrate each combination of modalities.
        """
        for i in range(1, len(self.estimators) + 1):
            for comb in combinations(range(len(self.estimators)), i):
                probas = predictions[:, np.array(comb)]
                if len(comb) == 1:
                    mask = (probas > 0).reshape(-1)
                else:
                    w = weights[:, np.array(comb)]
                    mask = np.any(probas > 0, axis=1).reshape(-1)
                    temp = np.sum(w, axis=1)
                    w[temp > 0] = w[temp > 0] / (temp[temp > 0].reshape(-1, 1))
                    probas = np.sum(probas * w, axis=1)

                self.fitted_meta_estimators_[comb] = LogisticRegression(
                    class_weight="balanced"
                ).fit(probas[mask].reshape(-1, 1), y[mask])
        return

    def predict_proba(self, X, estim_ind=None):
        """
        Late fusion probability estimates

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        estim_ind: tuple of integers.
            Tuple representing a multimodal combination (e.g. (i, j, k) corresponds to the combination of the ith, the
            jth and the kth estimators in self.fitted_estimators_). If None all the multimodal combination with all the
            fitted unimodal predictors is considered.

        Returns
        -------
        probas: array of shape (n_samples, 2).
            Probability of the samples for each class. If no modality are availbale for the sample, returns 0.5 for
            both classes.
        """
        fitted_estimators = (
            [self.fitted_estimators_[i] for i in estim_ind]
            if estim_ind is not None
            else self.fitted_estimators_
        )
        fitted_weights = (
            np.array([self.weights_[i] for i in estim_ind])
            if estim_ind is not None
            else self.weights_
        )

        # Collect predictions for each modality
        preds = np.zeros((X.shape[0], len(fitted_estimators)))
        weights = np.zeros((X.shape[0], len(fitted_weights)))
        for j, item in enumerate(fitted_estimators):
            Xpred = X[:, item[2]].copy()
            # bool_mask = ~(np.sum(np.isnan(Xpred), axis=1) > self.missing_threshold * len(item[2]))
            bool_mask = ~(
                np.sum(pd.isnull(Xpred), axis=1) > self.missing_threshold * len(item[2])
            )
            weights[:, j] = np.where(bool_mask, fitted_weights[j], 0)
            preds[bool_mask, j] = item[1].predict_proba(Xpred[bool_mask, :])[:, 1]

        # Calibrate the predictions and predict probas
        if self.calibration is not None:
            if self.calibration == "late":
                probas = self._predict_calibrate_late(preds, weights, estim_ind)
            elif self.calibration == "early":
                probas = self._predict_calibrate_early(preds, weights, estim_ind)
            else:
                raise ValueError(
                    "'early', 'late' or None are the only values available for calibration parameter"
                )
        else:
            probas = self._predict_uncalibrated(preds, weights)
        return np.hstack([1 - probas, probas])

    @staticmethod
    def _predict_uncalibrated(preds, weights):
        """
        Return weighted sum of available unimodal predictions
        """
        temp = np.sum(weights, axis=1)
        weights[temp > 0] = weights[temp > 0] / (temp[temp > 0].reshape(-1, 1))
        probas = np.sum(preds * weights, axis=1)
        return np.where(temp == 0, 0.5, probas).reshape(-1, 1)

    def _predict_calibrate_early(self, preds, weights, estim_ind):
        """
        Return weighted sum of available and calibrated unimodal predictions
        """
        temp = np.sum(weights, axis=1)
        weights[temp > 0] = weights[temp > 0] / (temp[temp > 0].reshape(-1, 1))
        list_meta_estimators = (
            [self.fitted_meta_estimators_[i] for i in estim_ind]
            if estim_ind is not None
            else list(self.fitted_meta_estimators_.values())
        )
        for j, meta in enumerate(list_meta_estimators):
            preds[:, j] = np.where(
                weights[:, j] != 0,
                meta.predict_proba(preds[:, j].reshape(-1, 1))[:, 1],
                0,
            )
        probas = np.sum(preds * weights, axis=1)
        return np.where(temp == 0, 0.5, probas).reshape(-1, 1)

    def _predict_calibrate_late(self, preds, weights, estim_ind):
        """
        Return calibrated weighted sum of availbale unimodal predictions
        """
        temp = np.sum(weights, axis=1)
        weights[temp > 0] = weights[temp > 0] / (temp[temp > 0].reshape(-1, 1))
        probas = np.sum(preds * weights, axis=1)
        meta_estimator = (
            self.fitted_meta_estimators_[estim_ind]
            if estim_ind is not None
            else list(self.fitted_meta_estimators_.values())[-1]
        )
        return np.where(
            temp == 0, 0.5, meta_estimator.predict_proba(probas.reshape(-1, 1))[:, 1]
        ).reshape(-1, 1)

    def find_logrank_threshold(
        self, X, ysurv, estim_ind, percentile_min=30, percentile_max=70
    ):
        """
        Find the best cutoff that optimize the stratification of samples with respect to survival data (using logrank
        test).

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        ysurv: structured array of shape (n_samples,) see sksurv.util.Surv (from scikit-survival)
            Structured array for survival data associated with X.

        estim_ind: tuple of integers.
            Tuple representing a multimodal combination (e.g. (i, j, k) corresponds to the combination of the ith, the
            jth and the kth estimators in self.fitted_estimators_). If None all the multimodal combination with all the
            fitted unimodal predictors is considered.

        percentile_min: int in [0, 100]
            Minimum value of the percentile range used to explore various cutoff values for predicted probabilities

        percentile_max: int in [0, 100]
            Maximum value of the percentile range used to explore various cutoff values for predicted probabilities

        Returns
        -------
        cutoff: float.
            Best cutoff for the predicted probabilities that otpimize the log-rank test.
        """
        risk_score = self.predict_proba(X, estim_ind=estim_ind)[:, 1]
        bool_mask = risk_score == 0.5
        cutoffs, pvals = [], []
        risk_score_new, y_new = risk_score[~bool_mask], ysurv[~bool_mask]
        for p in np.arange(percentile_min, percentile_max + 1):
            c = np.percentile(risk_score_new, p)
            group1 = risk_score_new <= c
            group2 = risk_score_new > c
            test = logrank_test(
                durations_A=y_new[group1]["time"],
                durations_B=y_new[group2]["time"],
                event_observed_A=1 * (y_new[group1]["event"]),
                event_observed_B=1 * (y_new[group2]["event"]),
            )
            cutoffs.append(c)
            pvals.append(test.summary["p"].values[0])
        return cutoffs[np.argmin(pvals)]


def _collect(Xdata, ydata, estimator, bmask, feat, train, test):
    Xtrain, Xtest, ytrain, ytest = (
        Xdata[np.intersect1d(np.where(bmask)[0], train), :],
        Xdata[np.intersect1d(np.where(bmask)[0], test), :],
        ydata[np.intersect1d(np.where(bmask)[0], train)],
        ydata[np.intersect1d(np.where(bmask)[0], test)],
    )
    tempbis = clone(estimator).fit(Xtrain[:, feat], ytrain)
    return (
        np.intersect1d(np.where(bmask)[0], test),
        tempbis.predict_proba(Xtest[:, feat])[:, 1],
    )


class LateFusionSurvival(BaseEstimator):
    """
    Late fusion survival model for multimodal integration.

    Parameters
    ----------
     estimators: list of (str, estimator, list, dict) tuples.
        List of unimodal estimators to fit and fuse. Each unimodal estimator is associate with a tuple
        (`name`, `estimator` , `features`, `tune_dict`) where `name` is a string and corresponds to the name of the
        estimator, `estimator` ia scikit-learn estimator inheriting from BaseEstimator (can be a Pipeline), `features`
        is a list of indexes corresponding to the columns of the data associated with the modality of intereset, and
        `tune_dict` is either a dictionnary or a tuple (dict, n_iterations) for hyperparameter tuning and GridSearch or
        RandomSearch strategy respectively.

    cv: cross-validation generator
        cross-validation scheme for hyperparameter tuning (if `tuning` is not None) and/or calibration (if `calibration`
        is not None). The default is None

    score: str or callable.
        Score to use for tuning the unimodal models or weighting them at the late fusion step (i.e. sum of the unimodal
        predictions weighted by the performance of each unimodal model estimated with cross-validation).
        See sklearn.model_selection.cross_val_score for more details. The default is None.

    random_score: float.
        Random score for classification. Used when weighting the unimodal models with their estimated score. Weights
        will be max(score - random_score, 0). Unimodal models whose estimated performance is below the random_score
        will not be taken into account. The default is 0.5

    sup_weights: bool.
        Whether to use weights associated with the cross-validation performance of each unimodal model. If false no
        weights are used when fusing the unimodal predictions. The default is False.

    missing_threshold: float in ]0, 1].
        Minimum frequency of missing values to consider a whole modality missing (e.g., if `missing_threshold = 0.9` it
        means that for each sample and each modality at least 90% of the features associated with this modality must be
        missing to consider the whole modality missing). The default is 0.9.

    tuning: str or None.
        Strategy for tuning each model. Either 'gridsearch' for GridSearchCV or 'randomsearch' for RandomSearchCV. If
        None no hyperparameter tuning will be performed. The default is None.

    n_jobs: int.
        Number of jobs to run in parallel for hyperparameter tuning, collecting the predictions for calibration, or
        estimating the performance of each unimodal model with cross-validation. The default is None.

    calibration: bool.
        If True each unimodal model is associated with a tuple (mean, std) estimated on predictions collected with
        cross-validation. The predictions of each unimodal model are then standardized before the late fusion step.

    Attributes
    ----------
    best_params_: list of dict or empty list.
        List of best parameters for each unimodal predictor (output of GridSearchCV or RandomSearchCV). It follows the
        same order as the one of `estimators` list. If `tuning` is None returns an empty list (i.e., no hyperparameter
        tuning is performed).

    weights_: list of float.
        List of the weights associated to each modality and used at the late fusion stage for weighted sum.

    fitted_estimators_: list of estimators.
        List of fitted unimodal estimators.
    """

    def __init__(
        self,
        estimators,
        cv,
        score=None,
        random_score=0.5,
        sup_weights=True,
        missing_threshold=0.9,
        tuning=None,
        n_jobs=None,
        calibration=True,
    ):
        self.estimators = estimators
        self.cv = cv
        self.score = score
        self.random_score = random_score
        self.sup_weights = sup_weights
        self.missing_threshold = missing_threshold
        self.tuning = tuning
        self.n_jobs = n_jobs
        self.calibration = calibration

        self.weights_ = []
        self.fitted_estimators_ = []
        self.best_params_ = []

    def _fit_estim(self, X, y, estim, features, grid, name):

        if (self.tuning is not None) and (len(grid) > 0):
            if self.tuning == "gridsearch":
                search = GridSearchCV(
                    estimator=clone(estim),
                    param_grid=grid,
                    cv=self.cv,
                    scoring=self.score,
                    n_jobs=self.n_jobs,
                )

            elif self.tuning == "randomsearch":
                search = RandomizedSearchCV(
                    estimator=clone(estim),
                    param_distributions=grid[1],
                    n_iter=grid[0],
                    scoring=self.score,
                    n_jobs=self.n_jobs,
                    cv=self.cv,
                )

            search.fit(X, y)

            if self.sup_weights:
                self.weights_.append(search.best_score_)
            else:
                self.weights_.append(1.0)

            temp = search.best_estimator_
            self.best_params_.append(search.best_params_)
            # print("Best params " + name + " :", search.best_params_)
            # print("Best score " + name + " :", search.best_score_)
        else:
            if self.sup_weights:
                self.weights_.append(
                    np.mean(
                        cross_val_score(
                            estimator=clone(estim),
                            X=X,
                            y=y,
                            cv=self.cv,
                            scoring=self.score,
                        )
                    )
                )
            else:
                self.weights_.append(1.0)
            temp = clone(estim).fit(X, y)

        # self.fitted_estimators_.append((name, temp, features))
        return temp

    def fit(self, X, y):
        """
        Fit the latefusion survival model.

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        y: structured array of shape (n_samples, ) see sksurv.util.Surv (from scikit-survival).
            Structured array for survival target/outcome

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        for name, estim, features, grid in self.estimators:
            Xnew = X[:, features]
            bool_mask = ~(
                np.sum(np.isnan(Xnew), axis=1) > self.missing_threshold * len(features)
            )
            Xnew, ynew = Xnew[bool_mask, :], y[bool_mask]

            fitted_estim = self._fit_estim(
                Xnew, ynew, estim=estim, features=features, grid=grid, name=name
            )
            if self.calibration:
                parallel = Parallel(n_jobs=self.n_jobs)
                collected_predictions = parallel(
                    delayed(_collect_surv)(
                        Xdata=X,
                        ydata=y,
                        estimator=estim,
                        bmask=bool_mask,
                        feat=features,
                        train=train,
                        test=test,
                    )
                    for train, test in self.cv.split(X, y)
                )
                temp = np.concatenate(collected_predictions)
                mean, std = np.mean(temp), np.std(temp)
            else:
                mean, std = None, None
            self.fitted_estimators_.append((name, fitted_estim, features, (mean, std)))

        self.weights_ = np.array(self.weights_) - self.random_score
        self.weights_ = np.where(self.weights_ > 0, self.weights_, 0)
        # if np.sum(self.weights_) > 0:
        #    self.weights_ = self.weights_/np.sum(self.weights_)
        return self

    def predict(self, X, estim_ind=None):
        """
        Predict risk scores

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        estim_ind: tuple of integers.
            Tuple representing a multimodal combination (e.g. (i, j, k) corresponds to the combination of the ith, the
            jth and the kth estimators in self.fitted_estimators_). If None all the multimodal combination with all the
            fitted unimodal predictors is considered.

        Returns
        -------
        risk_scores: array of shape (n_samples,).
            Predictied risk scores. If no modality are availbale for the sample, returns 0.
        """
        if estim_ind is not None:
            fitted_estimators = [self.fitted_estimators_[i] for i in estim_ind]
        else:
            fitted_estimators = self.fitted_estimators_

        preds = np.zeros((X.shape[0], len(fitted_estimators)))
        weights = np.zeros((X.shape[0], len(fitted_estimators)))
        for j, item in enumerate(fitted_estimators):
            Xpred = X[:, item[2]].copy()
            bool_mask = ~(
                np.sum(np.isnan(Xpred), axis=1) > self.missing_threshold * len(item[2])
            )
            weights[:, j] = np.where(bool_mask, self.weights_[j], 0)
            if self.calibration:
                mean = item[3][0]
                std = item[3][1] if item[3][1] != 0 else 1
                preds[bool_mask, j] = (
                    item[1].predict(Xpred[bool_mask, :]) - mean
                ) / std
            else:
                preds[bool_mask, j] = item[1].predict(Xpred[bool_mask, :])
        temp = np.sum(weights, axis=1)
        weights[temp > 0] = weights[temp > 0] / (temp[temp > 0].reshape(-1, 1))
        return np.sum(preds * weights, axis=1)

    def find_logrank_threshold(
        self, X, y, estim_ind, percentile_min=30, percentile_max=70
    ):
        """
        Find the best cutoff that optimize the stratification of samples with respect to survival data (using logrank
        test).

        Parameters
        ----------
        X: array of shape (n_samples, n_features)
            Multimodal array, concatenation of the features from the different modalities. Missing modalities are filled
            with NaNs values for each sample.

        y: structured array of shape (n_samples,) see sksurv.util.Surv (from scikit-survival)
            Structured array for survival data associated with X.

        estim_ind: tuple of integers.
            Tuple representing a multimodal combination (e.g. (i, j, k) corresponds to the combination of the ith, the
            jth and the kth estimators in self.fitted_estimators_). If None all the multimodal combination with all the
            fitted unimodal predictors is considered.

        percentile_min: int in [0, 100]
            Minimum value of the percentile range used to explore various cutoff values for predicted probabilities

        percentile_max: int in [0, 100]
            Maximum value of the percentile range used to explore various cutoff values for predicted probabilities

        Returns
        -------
        cutoff: float.
            Best cutoff for the predicted probabilities that otpimize the log-rank test.
        """
        risk_score = self.predict(X, estim_ind=estim_ind)
        bool_mask = risk_score == 0
        cutoffs, pvals = [], []
        risk_score_new, y_new = risk_score[~bool_mask], y[~bool_mask]
        for p in np.arange(percentile_min, percentile_max + 1):
            c = np.percentile(risk_score_new, p)
            group1 = risk_score_new <= c
            group2 = risk_score_new > c
            test = logrank_test(
                durations_A=y_new[group1]["time"],
                durations_B=y_new[group2]["time"],
                event_observed_A=1 * (y_new[group1]["event"]),
                event_observed_B=1 * (y_new[group2]["event"]),
            )
            cutoffs.append(c)
            pvals.append(test.summary["p"].values[0])
        return cutoffs[np.argmin(pvals)]


def _collect_surv(Xdata, ydata, estimator, bmask, feat, train, test):
    Xtrain, Xtest, ytrain, ytest = (
        Xdata[np.intersect1d(np.where(bmask)[0], train), :],
        Xdata[np.intersect1d(np.where(bmask)[0], test), :],
        ydata[np.intersect1d(np.where(bmask)[0], train)],
        ydata[np.intersect1d(np.where(bmask)[0], test)],
    )
    tempbis = clone(estimator).fit(Xtrain[:, feat], ytrain)
    return tempbis.predict(Xtest[:, feat])