# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/estimators.core.ipynb.

# %% ../../nbs/estimators.core.ipynb 3
from __future__ import annotations
import warnings
import itertools
import inspect
from abc import ABC, abstractmethod
from typing import List, Optional, Union, Callable, Dict, Tuple, Any, Sequence, Iterable

import pandas as pd
import numpy as np
import joblib

try:
    import ipywidgets
    from tqdm.notebook import tqdm
except ImportError:
    from tqdm import tqdm
from sklearn.base import ClassifierMixin, RegressorMixin, TransformerMixin, clone
from sklearn.model_selection._split import BaseCrossValidator, BaseShuffleSplit
from sklearn.model_selection import train_test_split
from sklearn.ensemble import (
    VotingClassifier,
    VotingRegressor,
    StackingClassifier,
    StackingRegressor,
)
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.dummy import DummyClassifier, DummyRegressor
from sklearn.model_selection import (
    cross_validate,
    cross_val_predict,
    GridSearchCV,
    RandomizedSearchCV,
)
from sklearn.exceptions import UndefinedMetricWarning

from ..utils.stats import cramers_v
from ..utils.hyperparameters import get_grid
from ..utils.estimate import get_target_info, element_to_list_maybe
from ..utils.utils import get_kwargs, non_default_repr
from ..plot import PoniardPlotFactory
from ..preprocessing import PoniardPreprocessor

# %% auto 0
__all__ = ['PoniardBaseEstimator']

# %% ../../nbs/estimators.core.ipynb 4
class PoniardBaseEstimator(ABC):
    """Base estimator that sets up all the functionality for the classifier and regressor.

    Parameters
    ----------
    estimators :
        Estimators to evaluate.
    metrics :
        Metrics to compute for each estimator. This is more restrictive than sklearn's scoring
        parameter, as it does not allow callable scorers. Single strings are cast to lists
        automatically.
    preprocess : bool, optional
        If True, impute missing values, standard scale numeric data and one-hot or ordinal
        encode categorical data.
    custom_preprocessor :
        Preprocessor used instead of the default preprocessing pipeline. It must be able to be
        included directly in a scikit-learn Pipeline.
    cv :
        Cross validation strategy. Either an integer, a scikit-learn cross validation object,
        or an iterable.
    verbose :
        Verbosity level. Propagated to every scikit-learn function and estimator.
    random_state :
        RNG. Propagated to every scikit-learn function and estimator. The default None sets
        random_state to 0 so that cross_validate results are comparable.
    n_jobs :
        Controls parallel processing. -1 uses all cores. Propagated to every scikit-learn
        function.
    plugins :
        Plugin instances that run in set moments of setup, fit and plotting.
    plot_options :
        :class:poniard.plot.plot_factory.PoniardPlotFactory instance specifying Plotly format
        options or None, which sets the default factory.
    """

    def __init__(
        self,
        estimators: Optional[
            Union[
                Sequence[ClassifierMixin],
                Dict[str, ClassifierMixin],
                Sequence[RegressorMixin],
                Dict[str, RegressorMixin],
            ]
        ] = None,
        metrics: Optional[Union[str, Dict[str, Callable], Sequence[str]]] = None,
        preprocess: bool = True,
        custom_preprocessor: Union[
            None, Pipeline, TransformerMixin, PoniardPreprocessor
        ] = None,
        cv: Union[int, BaseCrossValidator, BaseShuffleSplit, Sequence] = None,
        verbose: int = 0,
        random_state: Optional[int] = None,
        n_jobs: Optional[int] = None,
        plugins: Optional[Sequence[Any]] = None,
        plot_options: Optional[PoniardPlotFactory] = None,
    ):

        self._init_params = get_kwargs()
        # TODO: Ugly check that metrics conforms to expected types. Should improve.
        if metrics and (
            (
                isinstance(metrics, Sequence)
                and not all(isinstance(m, str) for m in metrics)
            )
            or (
                isinstance(metrics, Dict)
                and not all(isinstance(m, str) for m in metrics.keys())
                and not all(isinstance(m, Callable) for m in metrics.values())
            )
        ):
            raise ValueError(
                "metrics can only be a string, a sequence of strings, a dict with "
                "strings as keys and callables as values, or None."
            )
        self.metrics = metrics
        self.preprocess = preprocess
        self.custom_preprocessor = custom_preprocessor
        self.cv = cv
        self.verbose = verbose
        self.random_state = random_state or 0
        self.estimators = element_to_list_maybe(estimators)
        self.n_jobs = n_jobs

        self._memory = None

        self._init_plugins(plugins)
        self._init_plots(plot_options)

        self._added_estimators = {}
        self._removed_estimators = []

    def _init_plugins(self, plugins: Optional[Sequence[Any]] = None) -> None:
        self.plugins = element_to_list_maybe(plugins)
        if self.plugins:
            [setattr(plugin, "_poniard", self) for plugin in self.plugins]
        return

    def _init_plots(self, plot_options: Optional[PoniardPlotFactory] = None) -> None:
        self.plot_options = plot_options or PoniardPlotFactory()
        self.plot = self.plot_options
        self.plot._poniard = self
        return

    @property
    def poniard_task(self) -> Optional[str]:
        """Check whether self is a Poniard regressor or classifier.

        Returns
        -------
        Optional[str]
            "regression", "classification" or None
        """
        from poniard import PoniardRegressor, PoniardClassifier

        if isinstance(self, PoniardRegressor):
            return "regression"
        elif isinstance(self, PoniardClassifier):
            return "classification"
        else:
            return None

    def setup(
        self,
        X: Union[pd.DataFrame, np.ndarray, List],
        y: Union[pd.DataFrame, np.ndarray, List],
    ) -> PoniardBaseEstimator:
        """Acts as an orchestrator for Poniard estimators by setting up everything neeeded
        for `PoniardBaseEstimator.fit`.

        Converts inputs to arrays if necessary, sets `metrics`,
        `preprocessor`, `cv` and `pipelines`.

        After running `PoniardBaseEstimator.setup`, both `X` and `y` will be held as attributes.


        Parameters
        ----------
        X :
            Features
        y :
            Target.
        """
        self._run_plugin_method("on_setup_start")
        if not isinstance(X, (pd.DataFrame, pd.Series, np.ndarray)):
            X = np.array(X)
        if not isinstance(y, (pd.DataFrame, pd.Series, np.ndarray)):
            y = np.array(y)
        self.X = X
        self.y = y
        self._run_plugin_method("on_setup_data")

        self.target_info = get_target_info(self.y, self.poniard_task)
        print("Target info", "-----------", sep="\n")
        print(
            f"Type: {self.target_info['type_']}",
            f"Shape: {self.target_info['shape']}",
            f"Unique values: {self.target_info['nunique']}",
            sep="\n",
            end="\n\n",
        )
        if self.target_info["type_"] == "multiclass-multioutput":
            raise NotImplementedError(
                "multiclass-multioutput targets are not supported as "
                "no sklearn metrics support them."
            )

        if self.metrics:
            self.metrics = element_to_list_maybe(self.metrics)
        else:
            self.metrics = self._build_metrics()
        print(
            "Main metric",
            "-----------",
            self._first_scorer(sklearn_scorer=False),
            sep="\n",
            end="\n\n",
        )

        if self.preprocess:
            if self.custom_preprocessor and not isinstance(
                self.custom_preprocessor, PoniardPreprocessor
            ):
                self.preprocessor = self.custom_preprocessor
            else:
                self.preprocessor = self._build_preprocessor()
            self._pass_instance_attrs(self.preprocessor)

        self._run_plugin_method("on_setup_preprocessor")

        self.pipelines = self._build_pipelines()

        self.cv = self._build_cv()

        self._run_plugin_method("on_setup_end")
        return self

    def _build_preprocessor(
        self, assigned_types: Optional[Dict[str, List[Union[str, int]]]] = None
    ) -> Pipeline:
        """Build default preprocessor using `PoniardPreprocessor`.

        The preprocessor imputes missing values, scales numeric features and encodes categorical
        features according to inferred types.

        """
        if self.custom_preprocessor:
            self._poniard_preprocessor = self.custom_preprocessor
        else:
            self._poniard_preprocessor = PoniardPreprocessor()
        self._poniard_preprocessor._poniard = self
        self._memory = self._poniard_preprocessor._memory
        self._poniard_preprocessor.build()
        self.feature_types = self._poniard_preprocessor.feature_types
        return self._poniard_preprocessor.preprocessor

    @property
    @abstractmethod
    def _default_estimators(self) -> List[ClassifierMixin]:
        return []

    @property
    def estimators_(self):
        warnings.warn(
            "'estimators_' has been renamed to 'pipelines'",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.pipelines

    @property
    def preprocessor_(self):
        warnings.warn(
            "'preprocessor_' has been renamed to 'preprocessor'",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.preprocessor

    @property
    def metrics_(self):
        warnings.warn(
            "'metrics_' has been renamed to 'metrics'", DeprecationWarning, stacklevel=2
        )
        return self.metrics

    @property
    def cv_(self):
        warnings.warn(
            "'cv_' has been renamed to 'cv'", DeprecationWarning, stacklevel=2
        )
        return self.cv

    def show_results(
        self,
        std: bool = False,
        wrt_dummy: bool = False,
    ):
        warnings.warn(
            "'show_results' has been renamed to 'get_results'",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.get_results(return_train_scores=False, std=std, wrt_dummy=wrt_dummy)

    def _build_pipelines(
        self,
    ) -> Dict[str, Union[ClassifierMixin, RegressorMixin]]:
        """Build `pipelines` dict where keys are the estimator class names.

        Adds dummy estimators if not included during construction. Does nothing if
        `pipelines` exists.

        """
        if isinstance(self.estimators, dict):
            estimators = self.estimators.copy()
        elif self.estimators:
            estimators = {
                estimator.__class__.__name__: estimator for estimator in self.estimators
            }
        else:
            estimators = {
                estimator.__class__.__name__: estimator
                for estimator in self._default_estimators
            }
        # Take into account removed or added estimators.
        estimators.update(self._added_estimators)
        [estimators.pop(name) for name in self._removed_estimators]
        estimators = self._add_dummy_estimators(estimators)

        for estimator in estimators.values():
            self._pass_instance_attrs(estimator)

        pipelines = {}
        if self.preprocess:
            pipelines.update(
                {
                    name: Pipeline(
                        [("preprocessor", self.preprocessor), (name, estimator)],
                        memory=self._memory,
                    )
                    for name, estimator in estimators.items()
                }
            )
        else:
            pipelines.update(
                {
                    name: Pipeline([(name, estimator)])
                    for name, estimator in estimators.items()
                }
            )
        self._fitted_pipeline_ids = []
        return pipelines

    def _add_dummy_estimators(self, estimators: dict):
        if (
            "DummyClassifier" in estimators.keys()
            or "DummyRegressor" in estimators.keys()
        ):
            return estimators
        if self.poniard_task == "classification":
            estimators.update({"DummyClassifier": DummyClassifier(strategy="prior")})
        elif self.poniard_task == "regression":
            estimators.update({"DummyRegressor": DummyRegressor(strategy="mean")})
        return estimators

    @abstractmethod
    def _build_metrics(self) -> Union[Dict[str, Callable], List[str]]:
        """Build metrics."""
        return ["accuracy"]

    @abstractmethod
    def _build_cv(self):
        return self.cv

    def fit(self) -> PoniardBaseEstimator:
        """This is the main Poniard method. It uses scikit-learn's `cross_validate` function to
        score all `metrics` for every `pipelines`, using `cv` for cross validation.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        if not hasattr(self, "cv"):
            raise ValueError("`setup` must be called before `fit`.")
        self._run_plugin_method("on_fit_start")

        results = {}
        filtered_pipelines = {
            name: pipeline
            for name, pipeline in self.pipelines.items()
            if id(pipeline) not in self._fitted_pipeline_ids
        }
        pbar = tqdm(filtered_pipelines.items())
        for i, (name, pipeline) in enumerate(pbar):
            pbar.set_description(f"{name}")
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=UndefinedMetricWarning)
                warnings.filterwarnings(
                    "ignore", message=".*will be encoded as all zeros"
                )
                result = cross_validate(
                    pipeline,
                    self.X,
                    self.y,
                    scoring=self.metrics,
                    cv=self.cv,
                    return_train_score=True,
                    verbose=self.verbose,
                    n_jobs=self.n_jobs,
                )
            results.update({name: result})
            self._fitted_pipeline_ids.append(id(pipeline))
            if i == len(pbar) - 1:
                pbar.set_description("Completed")
        if hasattr(self, "_experiment_results"):
            self._experiment_results.update(results)
        else:
            self._experiment_results = results

        self._process_results()
        self._process_long_results()
        self._run_plugin_method("on_fit_end")
        return self

    def _predict(
        self, method: str, estimator_names: Optional[Sequence[str]] = None
    ) -> Dict[str, np.ndarray]:
        """Helper method for predicting targets or target probabilities with cross validation.
        Accepts predict, predict_proba, predict_log_proba or decision_function."""
        if not hasattr(self, "cv"):
            raise ValueError("`setup` must be called before `predict`.")
        X, y = self.X, self.y
        estimator_names = element_to_list_maybe(estimator_names)
        if not estimator_names:
            estimator_names = [estimator for estimator in self.pipelines.keys()]
        results = {}
        pbar = tqdm(estimator_names)
        for i, name in enumerate(pbar):
            pbar.set_description(f"{name}")
            pipeline = self.pipelines[name]
            try:
                result = cross_val_predict(
                    pipeline,
                    X,
                    y,
                    cv=self.cv,
                    method=method,
                    verbose=self.verbose,
                    n_jobs=self.n_jobs,
                )
            except AttributeError:
                warnings.warn(
                    f"{name} does not support `{method}` method. Filling with nan.",
                    stacklevel=2,
                )
                result = np.empty(self.y.shape)
                result[:] = np.nan
            results.update({name: result})

            if not hasattr(self, "_experiment_results"):
                self._experiment_results = {}
                self._experiment_results.update({name: {method: result}})
            elif name not in self._experiment_results:
                self._experiment_results.update({name: {method: result}})
            else:
                self._experiment_results[name][method] = result

            if i == len(pbar) - 1:
                pbar.set_description("Completed")
        return results

    def predict(
        self, estimator_names: Optional[Sequence[str]] = None
    ) -> Dict[str, np.ndarray]:
        """Get cross validated target predictions where each sample belongs to a single test set.

        Parameters
        ----------
        estimator_names :
            Estimators to include. If None, predict all estimators.

        Returns
        -------
        Dict
            Dict where keys are estimator names and values are numpy arrays of predictions.
        """
        return self._predict(method="predict", estimator_names=estimator_names)

    def predict_proba(
        self, estimator_names: Optional[Sequence[str]] = None
    ) -> Dict[str, np.ndarray]:
        """Get cross validated target probability predictions where each sample belongs to a
        single test set.

        Returns
        -------
        Dict
            Dict where keys are estimator names and values are numpy arrays of prediction
            probabilities.
        """
        return self._predict(method="predict_proba", estimator_names=estimator_names)

    def decision_function(
        self, estimator_names: Optional[Sequence[str]] = None
    ) -> Dict[str, np.ndarray]:
        """Get cross validated decision function predictions where each sample belongs to a
        single test set.

        Parameters
        ----------
        estimator_names :
            Estimators to include. If None, predict all estimators.

        Returns
        -------
        Dict
            Dict where keys are estimator names and values are numpy arrays of decision functions.
        """
        return self._predict(
            method="decision_function", estimator_names=estimator_names
        )

    def predict_all(
        self, estimator_names: Optional[Sequence[str]] = None
    ) -> Tuple[Dict[str, np.ndarray]]:
        """Get cross validated target predictions, probabilities and decision functions
        where each sample belongs to a test set.

        Parameters
        ----------
        estimator_names :
            Estimators to include. If None, predict all estimators.

        Returns
        -------
        Tuple[Dict]
            Tuple of dicts where keys are estimator names and values are numpy arrays of
            predictions.
        """
        return (
            self._predict(method="predict", estimator_names=estimator_names),
            self._predict(method="predict_proba", estimator_names=estimator_names),
            self._predict(method="decision_function", estimator_names=estimator_names),
        )

    def get_results(
        self,
        return_train_scores: bool = False,
        std: bool = False,
        wrt_dummy: bool = False,
    ) -> Union[Tuple[pd.DataFrame, pd.DataFrame], pd.DataFrame]:
        """Return dataframe containing scoring results. By default returns the mean score and fit
        and score times. Optionally returns standard deviations as well.

        Parameters
        ----------
        return_train_scores :
            If False, only return test scores.
        std :
            Whether to return standard deviation of the scores. Default False.
        wrt_dummy :
            Whether to compute each score/time with respect to the dummy estimator results. Default
            False.

        Returns
        -------
        Union[Tuple[pd.DataFrame, pd.DataFrame], pd.DataFrame]
            Results
        """
        means = self._means
        stds = self._stds
        if not return_train_scores:
            means = means.loc[
                :, means.columns.str.contains("test_|fit|score", regex=True)
            ]
            stds = stds.loc[:, stds.columns.str.contains("test_|fit|score", regex=True)]
        if wrt_dummy:
            dummy_means = means.loc[means.index.str.contains("Dummy")]
            dummy_stds = stds.loc[stds.index.str.contains("Dummy")]
            means = means / dummy_means.squeeze()
            stds = stds / dummy_stds.squeeze()
        if std:
            return means, stds
        else:
            return means

    def reassign_types(
        self,
        numeric: Optional[List[Union[str, int]]] = None,
        categorical_high: Optional[List[Union[str, int]]] = None,
        categorical_low: Optional[List[Union[str, int]]] = None,
        datetime: Optional[List[Union[str, int]]] = None,
        keep_remainder: bool = True,
    ) -> PoniardBaseEstimator:
        """Reassign feature types. By default, leaves ommitted features as they were.

        Parameters
        ----------
        numeric :
            List of column names or indices. Default None.
        categorical_high :
            List of column names or indices. Default None.
        categorical_low :
            List of column names or indices. Default None.
        datetime :
            List of column names or indices. Default None.
        keep_remainder :
            Whether to keep features not specified in the method parameters
            as is or drop them

        Returns
        -------
        PoniardBaseEstimator
            self.
        """
        numeric = numeric or []
        categorical_high = categorical_high or []
        categorical_low = categorical_low or []
        datetime = datetime or []
        if keep_remainder:
            assigned_types = self.feature_types.copy()
            swapped = numeric + categorical_high + categorical_low + datetime
            for k in self.feature_types.keys():
                assigned_types[k] = [x for x in assigned_types[k] if x not in swapped]
            for k, new in zip(
                assigned_types.keys(),
                [numeric, categorical_high, categorical_low, datetime],
            ):
                assigned_types[k] = assigned_types[k] + new
        else:
            assigned_types = {
                "numeric": numeric or [],
                "categorical_high": categorical_high or [],
                "categorical_low": categorical_low or [],
                "datetime": datetime or [],
            }
        print("Assigned feature types", "----------------------", sep="\n")
        assigned_types_df = pd.DataFrame.from_dict(
            assigned_types, orient="index"
        ).T.fillna("")
        try:
            # Try to print the table nicely
            from IPython.display import display, HTML

            display(HTML(assigned_types_df.to_html()))
            print("\n")
        except ImportError:
            print(assigned_types_df)

        self.feature_types = assigned_types
        self._poniard_preprocessor.feature_types = assigned_types
        self._poniard_preprocessor.build()
        self.preprocessor = self._poniard_preprocessor.preprocessor
        self._run_plugin_method("on_reassign_types")
        self.pipelines = self._build_pipelines()
        return self

    def add_preprocessing_step(
        self,
        step: Union[
            Union[Pipeline, TransformerMixin, ColumnTransformer],
            Tuple[str, Union[Pipeline, TransformerMixin, ColumnTransformer]],
        ],
        position: Union[str, int] = "end",
    ) -> Pipeline:
        """Add a preprocessing step.

        Parameters
        ----------
        step :
            A tuple of (str, transformer) or a scikit-learn transformer. Note that
            the transformer can also be a Pipeline or ColumnTransformer.
        position :
            Either an integer denoting before which step in the existing preprocessing pipeline
            the new step should be added, or 'start' or 'end'.

        Returns
        -------
        PoniardPreprocessor
            self
        """
        if not isinstance(position, int) and position not in ["start", "end"]:
            raise ValueError("`position` can only be int, 'start' or 'end'.")
        existing_preprocessor = self.preprocessor
        if not isinstance(step, Tuple):
            step = (f"step_{step.__class__.__name__.lower()}", step)
        if isinstance(position, str) and isinstance(existing_preprocessor, Pipeline):
            if position == "start":
                position = 0
            elif position == "end":
                position = len(existing_preprocessor.steps)
        if isinstance(existing_preprocessor, Pipeline):
            existing_preprocessor.steps.insert(position, step)
        else:
            if isinstance(position, int):
                raise ValueError(
                    "If the existing preprocessor is not a Pipeline, only 'start' and "
                    "'end' are accepted as `position`."
                )
            if position == "start":
                self.preprocessor = Pipeline(
                    [step, ("initial_preprocessor", self.preprocessor)],
                    memory=self._memory,
                )
            else:
                self.preprocessor = Pipeline(
                    [("initial_preprocessor", self.preprocessor), step],
                    memory=self._memory,
                )
        self.pipelines = self._build_pipelines()
        self._run_plugin_method("on_add_preprocessing_step")
        return self

    def add_estimators(
        self, estimators: Union[Dict[str, ClassifierMixin], Sequence[ClassifierMixin]]
    ) -> PoniardBaseEstimator:
        """Include new estimator. This is the recommended way of adding an estimator (as opposed
        to modifying `pipelines` directly), since it also injects random state, n_jobs
        and verbosity.

        Parameters
        ----------
        estimators :
            Estimators to add.

        Returns
        -------
        PoniardBaseEstimator
            Self.

        """
        estimators = element_to_list_maybe(estimators)
        if not isinstance(estimators, dict):
            new_estimators = {
                estimator.__class__.__name__: estimator for estimator in estimators
            }
        else:
            new_estimators = estimators
        for new_estimator in new_estimators.values():
            self._pass_instance_attrs(new_estimator)
        self._added_estimators.update(new_estimators)
        if self.preprocess:
            self.pipelines.update(
                {
                    name: Pipeline(
                        [("preprocessor", self.preprocessor), (name, estimator)],
                        memory=self._memory,
                    )
                    for name, estimator in new_estimators.items()
                }
            )
        else:
            self.pipelines.update(
                {
                    name: Pipeline([(name, estimator)])
                    for name, estimator in new_estimators.items()
                }
            )
        self._run_plugin_method("on_add_estimators")
        return self

    def remove_estimators(
        self, estimator_names: Sequence[str], drop_results: bool = True
    ) -> PoniardBaseEstimator:
        """Remove estimators. This is the recommended way of removing an estimator (as opposed
        to modifying `pipelines` directly), since it also removes the associated rows from
        the results tables.

        Parameters
        ----------
        estimator_names :
            Estimators to remove.
        drop_results :
            Whether to remove the results associated with the estimators. Default True.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        estimator_names = element_to_list_maybe(estimator_names)
        self._removed_estimators.extend(estimator_names)
        pruned_estimators = {
            k: v for k, v in self.pipelines.items() if k not in estimator_names
        }
        if len(pruned_estimators) == 0:
            raise ValueError("Cannot remove all estimators.")
        self.pipelines = pruned_estimators
        if drop_results and hasattr(self, "_means"):
            self._means = self._means.loc[~self._means.index.isin(estimator_names)]
            self._stds = self._stds.loc[~self._stds.index.isin(estimator_names)]
            self._experiment_results = {
                k: v
                for k, v in self._experiment_results.items()
                if k not in estimator_names
            }
            self._process_long_results()
        self._run_plugin_method("on_remove_estimators")
        return self

    def get_estimator(
        self,
        estimator_name: str,
        include_preprocessor: bool = True,
        retrain: bool = False,
    ) -> Union[Pipeline, ClassifierMixin, RegressorMixin]:
        """Obtain an estimator in `pipelines` by name. This is useful for extracting default
        estimators or hyperparmeter-optimized estimators (after using
        `PoniardBaseEstimator.tune_estimator`).

        Parameters
        ----------
        estimator_name :
            Estimator name.
        include_preprocessor :
            Whether to return a pipeline with a preprocessor or just the estimator. Default True.
        retrain :
            Whether to retrain with full data. Default False.

        Returns
        -------
        ClassifierMixin
            Estimator.
        """
        model = self.pipelines[estimator_name]
        if not include_preprocessor:
            model = model._final_estimator
        model = clone(model)
        if retrain:
            model.fit(self.X, self.y)
        self._run_plugin_method(
            "on_get_estimator", estimator=model, name=estimator_name
        )
        return model

    def build_ensemble(
        self,
        method: str = "stacking",
        estimator_names: Optional[Sequence[str]] = None,
        top_n: Optional[int] = 3,
        sort_by: Optional[str] = None,
        ensemble_name: Optional[str] = None,
        **kwargs,
    ) -> PoniardBaseEstimator:
        """Combine estimators into an ensemble.

        By default, orders estimators according to the first metric.

        Parameters
        ----------
        method :
            Ensemble method. Either "stacking" or "voting". Default "stacking".
        estimator_names :
            Names of estimators to include. Default None, which uses `top_n`
        top_n :
            How many of the best estimators to include.
        sort_by :
            Which metric to consider for ordering results. Default None, which uses the first metric.
        ensemble_name :
            Ensemble name when adding to `pipelines`. Default None.
        kwargs :
            Passed to the ensemble class constructor.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        if method not in ["voting", "stacking"]:
            raise ValueError("Method must be either voting or stacking.")
        estimator_names = element_to_list_maybe(estimator_names)
        if estimator_names:
            models = [
                (name, self.pipelines[name]._final_estimator)
                for name in estimator_names
            ]
        else:
            if sort_by:
                sorter = sort_by
            else:
                sorter = self._means.columns[0]
            models = [
                (name, self.pipelines[name]._final_estimator)
                for name in self._means.sort_values(sorter, ascending=False).index[
                    :top_n
                ]
            ]
        if method == "voting":
            if self.poniard_task == "classification":
                ensemble = VotingClassifier(
                    estimators=models, verbose=self.verbose, **kwargs
                )
            else:
                ensemble = VotingRegressor(
                    estimators=models, verbose=self.verbose, **kwargs
                )
        else:
            if self.poniard_task == "classification":
                ensemble = StackingClassifier(
                    estimators=models, verbose=self.verbose, cv=self.cv, **kwargs
                )
            else:
                ensemble = StackingRegressor(
                    estimators=models, verbose=self.verbose, cv=self.cv, **kwargs
                )
        ensemble_name = ensemble_name or ensemble.__class__.__name__
        self.add_estimators(estimators={ensemble_name: ensemble})
        return self

    def get_predictions_similarity(
        self,
        on_errors: bool = True,
    ) -> pd.DataFrame:
        """Compute correlation/association between cross validated predictions for each estimator.

        This can be useful for ensembling.

        Parameters
        ----------
        on_errors :
            Whether to compute similarity on prediction errors instead of predictions. Default
            True.

        Returns
        -------
        pd.DataFrame
            Similarity.
        """
        if self.y.ndim > 1:
            raise ValueError("y must be a 1-dimensional array.")
        raw_results = {
            name: self._get_or_compute_prediction(estimator_name=name, method="predict")
            for name in self.pipelines.keys()
        }
        results = raw_results.copy()
        for name, result in raw_results.items():
            if on_errors:
                if self.poniard_task == "regression":
                    results[name] = self.y - result
                else:
                    results[name] = np.where(result == self.y, 1, 0)
        results = pd.DataFrame(results)
        if self.poniard_task == "classification":
            estimator_names = [x for x in results.columns if x != "DummyClassifier"]
            table = pd.DataFrame(
                data=np.nan, index=estimator_names, columns=estimator_names
            )
            for row, col in itertools.combinations_with_replacement(
                table.index[::-1], 2
            ):
                cramer = cramers_v(results[row], results[col])
                if row == col:
                    table.loc[row, col] = 1
                else:
                    table.loc[row, col] = cramer
                    table.loc[col, row] = cramer
        else:
            table = results.drop("DummyRegressor", axis=1).corr()
        return table

    def tune_estimator(
        self,
        estimator_name: str,
        grid: Optional[Dict] = None,
        mode: str = "grid",
        tuned_estimator_name: Optional[str] = None,
        **kwargs,
    ) -> Union[GridSearchCV, RandomizedSearchCV]:
        """Hyperparameter tuning for a single estimator.

        Parameters
        ----------
        estimator_name :
            Estimator to tune.
        grid :
            Hyperparameter grid. Default None, which uses the grids available for default
            estimators.
        mode :
            Type of search. Eitherr "grid", "halving" or "random". Default "grid".
        tuned_estimator_name :
            Estimator name when adding to `pipelines`. Default None.
        kwargs :
            Passed to the search class constructor.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        X, y = self.X, self.y
        estimator = clone(self.pipelines[estimator_name])
        if not grid:
            try:
                grid = get_grid(estimator_name)
                grid = {f"{estimator_name}__{k}": v for k, v in grid.items()}
            except KeyError:
                raise NotImplementedError(
                    f"Estimator {estimator_name} has no predefined hyperparameter grid, so it has to be supplied."
                )
        self._pass_instance_attrs(estimator)

        scoring = self._first_scorer(sklearn_scorer=True)
        if mode == "random":
            search = RandomizedSearchCV(
                estimator,
                grid,
                scoring=scoring,
                cv=self.cv,
                verbose=self.verbose,
                n_jobs=self.n_jobs,
                random_state=self.random_state,
                **kwargs,
            )
        elif mode == "halving":
            from sklearn.experimental import enable_halving_search_cv
            from sklearn.model_selection import HalvingGridSearchCV

            search = HalvingGridSearchCV(
                estimator,
                grid,
                scoring=scoring,
                cv=self.cv,
                verbose=self.verbose,
                n_jobs=self.n_jobs,
                random_state=self.random_state,
                **kwargs,
            )
        else:
            search = GridSearchCV(
                estimator,
                grid,
                scoring=scoring,
                cv=self.cv,
                verbose=self.verbose,
                n_jobs=self.n_jobs,
                **kwargs,
            )
        search.fit(X, y)
        tuned_estimator_name = tuned_estimator_name or f"{estimator_name}_tuned"
        self.add_estimators(
            estimators={
                tuned_estimator_name: clone(search.best_estimator_._final_estimator)
            }
        )
        return self

    def _process_results(self) -> None:
        """Compute mean and standard deviations of  experiment results."""
        # TODO: This processes every result, even those that were processed
        # in previous runs (before add_estimators). Should be made more efficient
        results = pd.DataFrame(self._experiment_results).T
        results = results.loc[
            :,
            [
                x
                for x in results.columns
                if x not in ["predict", "predict_proba", "decision_function"]
            ],
        ]
        means = results.apply(lambda x: np.mean(x.values.tolist(), axis=1))
        stds = results.apply(lambda x: np.std(x.values.tolist(), axis=1))
        means = means[list(means.columns[2:]) + ["fit_time", "score_time"]]
        stds = stds[list(stds.columns[2:]) + ["fit_time", "score_time"]]
        self._means = means.sort_values(means.columns[0], ascending=False)
        self._stds = stds.reindex(self._means.index)
        return

    def _process_long_results(self) -> None:
        """Prepare experiment results for plotting."""
        base = pd.DataFrame(self._experiment_results).T
        melted = (
            base.rename_axis("Model")
            .reset_index()
            .melt(id_vars="Model", var_name="Metric", value_name="Score")
            .explode("Score")
        )
        melted["Type"] = "Fold"
        means = melted.groupby(["Model", "Metric"])["Score"].mean().reset_index()
        means["Type"] = "Mean"
        melted = pd.concat([melted, means])
        melted["Model"] = melted["Model"].str.replace(
            "Classifier|Regressor", "", regex=True
        )

        self._long_results = melted
        return

    def _first_scorer(self, sklearn_scorer: bool) -> Union[str, Callable]:
        """Helper method to get the first scoring function or name."""
        if isinstance(self.metrics, Sequence):
            return self.metrics[0]
        elif isinstance(self.metrics, dict):
            if sklearn_scorer:
                return list(self.metrics.values())[0]
            else:
                return list(self.metrics.keys())[0]
        else:
            raise ValueError(
                "self.metrics can only be a sequence of str or dict of str: callable."
            )

    def _train_test_split_from_cv(self):
        """Split data in a 80/20 fashion following the cross-validation strategy defined in the constructor."""
        if isinstance(self.cv, (int, Iterable)):
            cv_params_for_split = {}
        else:
            cv_params_for_split = {
                k: v
                for k, v in vars(self.cv).items()
                if k in ["shuffle", "random_state"]
            }
            stratify = self.y if "Stratified" in self.cv.__class__.__name__ else None
            cv_params_for_split.update({"stratify": stratify})
        return train_test_split(self.X, self.y, test_size=0.2, **cv_params_for_split)

    def _pass_instance_attrs(self, obj: Union[ClassifierMixin, RegressorMixin]):
        """Helper method to propagate instance attributes to objects."""
        for attr, value in zip(
            ["random_state", "verbose", "verbosity"],
            [self.random_state, self.verbose, self.verbose],
        ):
            if hasattr(obj, attr):
                setattr(obj, attr, value)
        return

    def _run_plugin_method(self, method: str, **kwargs):
        """Helper method to run plugin methods by name."""
        if not self.plugins:
            return
        for plugin in self.plugins:
            fetched_method = getattr(plugin, method, None)
            if callable(fetched_method):
                accepted_kwargs = inspect.getargs(fetched_method.__code__).args
                matched_kwargs = {
                    k: v for k, v in kwargs.items() if k in accepted_kwargs
                }
                fetched_method(**matched_kwargs)
        return

    def _get_or_compute_prediction(self, estimator_name: str, method: str):
        """Get predictions (either predict, predict_proba or decision_function) for a given
        estimator or compute if not available."""
        try:
            return self._experiment_results[estimator_name][method]
        except KeyError:
            self._predict(method=method, estimator_names=[estimator_name])
            return self._experiment_results[estimator_name][method]

    def __repr__(self):
        return non_default_repr(self)

    def __add__(
        self,
        estimators: Union[
            Dict[str, Union[ClassifierMixin, RegressorMixin]],
            Sequence[Union[ClassifierMixin, RegressorMixin]],
        ],
    ) -> PoniardBaseEstimator:
        """Add estimators to a Poniard Estimator.

        Parameters
        ----------
        estimators :
            List or dict of estimators to add.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        estimators = element_to_list_maybe(estimators)
        return self.add_estimators(estimators)

    def __sub__(self, estimator_names: Sequence[str]) -> PoniardBaseEstimator:
        """Remove an estimator and its results.

        Parameters
        ----------
        estimator :
            List of estimators names.

        Returns
        -------
        PoniardBaseEstimator
            Self.
        """
        estimator_names = element_to_list_maybe(estimator_names)
        return self.remove_estimators(estimator_names, drop_results=True)

    def __getitem__(
        self, estimator_name: str
    ) -> Union[Pipeline, ClassifierMixin, RegressorMixin]:
        """Get an estimator by indexing with its name

        Parameters
        ----------
        estimator_name :
            Estimator name as string.

        Returns
        -------
        Union[Pipeline, ClassifierMixin, RegressorMixin]
            Built estimator.
        """
        return self.get_estimator(estimator_name)
