import copy
import logging

import pandas as pd

from .dataset import TabularDataset
from ..base.base_predictor import BasePredictor
from ...utils import plot_performance_vs_trials, plot_summary_of_models, plot_tabular_models, verbosity2loglevel
from ...utils.tabular.ml.constants import REGRESSION
from ...utils.tabular.ml.learner.abstract_learner import AbstractLearner as Learner  # TODO: Keep track of true type of learner for loading
from ...utils.tabular.ml.trainer.abstract_trainer import AbstractTrainer  # TODO: Keep track of true type of trainer for loading
from ...utils.tabular.ml.utils import setup_outputdir

__all__ = ['TabularPredictor']

logger = logging.getLogger()  # return root logger


class TabularPredictor(BasePredictor):
    """ Object returned by `fit()` in Tabular Prediction tasks.
        Use for making predictions on new data and viewing information about models trained during `fit()`.

        Attributes
        ----------
        output_directory : str
            Path to directory where all models used by this Predictor are stored.
        problem_type : str
            What type of prediction problem this Predictor has been trained for.
        eval_metric : function or str
            What metric is used to evaluate predictive performance.
        label_column : str
            Name of table column that contains data from the variable to predict (often referred to as: labels, response variable, target variable, dependent variable, Y, etc).
        feature_types : dict
            Inferred data type of each predictive variable (i.e. column of training data table used to predict `label_column`).
        model_names : list
            List of model names trained during `fit()`.
        model_performance : dict
            Maps names of trained models to their predictive performance values attained on the validation dataset during `fit()`.
        class_labels : list
            For multiclass problems, this list contains the class labels in sorted order of `predict_proba()` output.
            For binary problems, this list contains the class labels in sorted order of `predict_proba(as_multiclass=True)` output.
                `class_labels[0]` corresponds to internal label = 0 (negative class), `class_labels[1]` corresponds to internal label = 1 (positive class).
                This is relevant for certain metrics such as F1 where True and False labels impact the metric score differently.
            For other problem types, will equal None.
            For example if `pred = predict_proba(x, as_multiclass=True)`, then ith index of `pred` provides predicted probability that `x` belongs to class given by `class_labels[i]`.
        class_labels_internal : list
            For multiclass problems, this list contains the internal class labels in sorted order of internal `predict_proba()` output.
            For binary problems, this list contains the internal class labels in sorted order of internal `predict_proba(as_multiclass=True)` output.
                The value will always be `class_labels_internal=[0, 1]` for binary problems, with 0 as the negative class, and 1 as the positive class.
            For other problem types, will equal None.
        class_labels_internal_map : dict
            For binary and multiclass classification problems, this dictionary contains the mapping of the original labels to the internal labels.
            For example, in binary classification, label values of 'True' and 'False' will be mapped to the internal representation `1` and `0`.
                Therefore, class_labels_internal_map would equal {'True': 1, 'False': 0}
            For other problem types, will equal None.
            For multiclass, it is possible for not all of the label values to have a mapping.
                This indicates that the internal models will never predict those missing labels, and training rows associated with the missing labels were dropped.

        Examples
        --------
        >>> from autogluon import TabularPrediction as task
        >>> train_data = task.Dataset(file_path='https://autogluon.s3-us-west-2.amazonaws.com/datasets/Inc/train.csv')
        >>> predictor = task.fit(train_data=train_data, label='class')
        >>> results = predictor.fit_summary()
        >>> test_data = task.Dataset(file_path='https://autogluon.s3-us-west-2.amazonaws.com/datasets/Inc/test.csv')
        >>> perf = predictor.evaluate(test_data)

    """

    def __init__(self, learner):
        """ Creates TabularPredictor object.
            You should not construct a TabularPredictor yourself, it is only intended to be produced during fit().

            Parameters
            ----------
            learner : `AbstractLearner` object
                Object that implements the `AbstractLearner` APIs.

            To access any learner method `func()` from this Predictor, use: `predictor._learner.func()`.
            To access any trainer method `func()` from this `Predictor`, use: `predictor._trainer.func()`.
        """
        self._learner: Learner = learner  # Learner object
        self._learner.persist_trainer(low_memory=True)
        self._trainer: AbstractTrainer = self._learner.load_trainer()  # Trainer object
        self.output_directory = self._learner.path
        self.problem_type = self._learner.problem_type
        self.eval_metric = self._learner.objective_func
        self.label_column = self._learner.label
        self.feature_types = self._trainer.feature_types_metadata
        self.model_names = self._trainer.get_model_names_all()  # TODO: Consider making this a function instead of a variable: This should never be de-synced with the output of self._trainer.get_model_names_all()
        self.model_performance = self._trainer.model_performance
        self.class_labels = self._learner.class_labels
        self.class_labels_internal = self._learner.label_cleaner.ordered_class_labels_transformed
        self.class_labels_internal_map = self._learner.label_cleaner.inv_map

    def predict(self, dataset, model=None, as_pandas=False, use_pred_cache=False, add_to_pred_cache=False):
        """ Use trained models to produce predicted labels (in classification) or response values (in regression).

            Parameters
            ----------
            dataset : str or :class:`TabularDataset` or `pandas.DataFrame`
                The dataset to make predictions for. Should contain same column names as training Dataset and follow same format
                (may contain extra columns that won't be used by Predictor, including the label-column itself).
                If str is passed, `dataset` will be loaded using the str value as the file path.
            model : str (optional)
                The name of the model to get predictions from. Defaults to None, which uses the highest scoring model on the validation set.
                Valid models are listed in this `predictor` by calling `predictor.model_names`
            as_pandas : bool (optional)
                Whether to return the output as a pandas Series (True) or numpy array (False)
            use_pred_cache : bool (optional)
                Whether to used previously-cached predictions for table rows we have already predicted on before
                (can speedup repeated runs of `predict()` on multiple datasets with overlapping rows between them).
            add_to_pred_cache : bool (optional)
                Whether these predictions should be cached for reuse in future `predict()` calls on the same table rows
                (can speedup repeated runs of `predict()` on multiple datasets with overlapping rows between them).

            Returns
            -------
            Array of predictions, one corresponding to each row in given dataset. Either numpy Ndarray or pandas Series depending on `as_pandas` argument.

        """
        dataset = self.__get_dataset(dataset)
        return self._learner.predict(X_test=dataset, model=model, as_pandas=as_pandas, use_pred_cache=use_pred_cache, add_to_pred_cache=add_to_pred_cache)

    def predict_proba(self, dataset, model=None, as_pandas=False, as_multiclass=False):
        """ Use trained models to produce predicted class probabilities rather than class-labels (if task is classification).

            Parameters
            ----------
            dataset : str or :class:`TabularDataset` or `pandas.DataFrame`
                The dataset to make predictions for. Should contain same column names as training Dataset and follow same format
                (may contain extra columns that won't be used by Predictor, including the label-column itself).
                If str is passed, `dataset` will be loaded using the str value as the file path.
            model : str (optional)
                The name of the model to get prediction probabilities from. Defaults to None, which uses the highest scoring model on the validation set.
                Valid models are listed in this `predictor` by calling `predictor.model_names`
            as_pandas : bool (optional)
                Whether to return the output as a pandas object (True) or numpy array (False).
                Pandas object is a DataFrame if this is a multiclass problem or `as_multiclass=True`, otherwise it is a Series.
                If the output is a DataFrame, the column order will be equivalent to `predictor.class_labels`.
            as_multiclass : bool (optional)
                Whether to return binary classification probabilities as if they were for multiclass classification.
                    Output will contain two columns, and if `as_pandas=True`, the column names will correspond to the binary class labels.
                    The columns will be the same order as `predictor.class_labels`.
                Only impacts output for binary classification problems.

            Returns
            -------
            Array of predicted class-probabilities, corresponding to each row in the given dataset.
            May be a numpy ndarray or pandas Series/DataFrame depending on `as_pandas` and `as_multiclass` arguments and the type of prediction problem.
            For binary classification problems, the output contains for each datapoint only the predicted probability of the positive class, unless you specify `as_multiclass=True`.
        """
        dataset = self.__get_dataset(dataset)
        return self._learner.predict_proba(X_test=dataset, model=model, as_pandas=as_pandas, as_multiclass=as_multiclass)

    def evaluate(self, dataset, silent=False):
        """ Report the predictive performance evaluated for a given Dataset.
            This is basically a shortcut for: `pred = predict(dataset); evaluate_predictions(dataset[label_column], preds, auxiliary_metrics=False)`
            that automatically uses `predict_proba()` instead of `predict()` when appropriate.

            Parameters
            ----------
            dataset : str or :class:`TabularDataset` or `pandas.DataFrame`
                This Dataset must also contain the label-column with the same column-name as specified during `fit()`.
                If str is passed, `dataset` will be loaded using the str value as the file path.

            silent : bool (optional)
                Should performance results be printed?

            Returns
            -------
            Predictive performance value on the given dataset, based on the `eval_metric` used by this Predictor.
        """
        dataset = self.__get_dataset(dataset)
        perf = self._learner.score(dataset)
        sign = self._learner.objective_func._sign
        perf = perf * sign  # flip negative once again back to positive (so higher is no longer necessarily better)
        if not silent:
            print("Predictive performance on given dataset: %s = %s" % (self.eval_metric, perf))
        return perf

    def evaluate_predictions(self, y_true, y_pred, silent=False, auxiliary_metrics=False, detailed_report=True):
        """ Evaluate the provided predictions against ground truth labels.
            Evaluation is based on the `eval_metric` previously specifed to `fit()`, or default metrics if none was specified.

            Parameters
            ----------
            y_true : list or `numpy.array`
                The ordered collection of ground-truth labels.
            y_pred : list or `numpy.array`
                The ordered collection of predictions.
                Caution: For certain types of `eval_metric` (such as 'roc_auc'), `y_pred` must be predicted-probabilities rather than predicted labels.
            silent : bool (optional)
                Should performance results be printed?
            auxiliary_metrics: bool (optional)
                Should we compute other (`problem_type` specific) metrics in addition to the default metric?
            detailed_report : bool (optional)
                Should we computed more detailed versions of the `auxiliary_metrics`? (requires `auxiliary_metrics = True`)

            Returns
            -------
            Scalar performance value if `auxiliary_metrics = False`.
            If `auxiliary_metrics = True`, returns dict where keys = metrics, values = performance along each metric.
        """
        return self._learner.evaluate(y_true=y_true, y_pred=y_pred, silent=silent,
                                      auxiliary_metrics=auxiliary_metrics, detailed_report=detailed_report)

    def leaderboard(self, dataset=None, only_pareto_frontier=False, silent=False):
        """
            Output summary of information about models produced during fit() as a pandas DataFrame.
            Includes information on test and validation scores for all models, model training times, inference times, and stack levels.
            Output DataFrame columns include:
                'model': The name of the model.
                'score_val': The validation score of the model on the 'eval_metric'.
                'pred_time_val': The inference time required to compute predictions on the validation data end-to-end.
                    Equivalent to the sum of all 'pred_time_val_marginal' values for the model and all of its base models.
                'fit_time': The fit time required to train the model end-to-end (Including base models if the model is a stack ensemble).
                    Equivalent to the sum of all 'fit_time_marginal' values for the model and all of its base models.
                'pred_time_val_marginal': The inference time required to compute predictions on the validation data (Ignoring inference times for base models).
                    Note that this ignores the time required to load the model into memory when bagging is disabled.
                'fit_time_marginal': The fit time required to train the model (Ignoring base models).
                'stack_level': The stack level of the model.
                    A model with stack level N can take any set of models with stack level less than N as input, with stack level 0 models having no model inputs.

            Parameters
            ----------
            dataset : str or :class:`TabularDataset` or `pandas.DataFrame` (optional)
                This Dataset must also contain the label-column with the same column-name as specified during fit().
                If specified, then the leaderboard returned will contain additional columns 'score_test', 'pred_time_test', and 'pred_time_test_marginal'.
                    'score_test': The score of the model on the 'eval_metric' for the dataset provided.
                    'pred_time_test': The true end-to-end wall-clock inference time of the model for the dataset provided.
                        Equivalent to the sum of all 'pred_time_test_marginal' values for the model and all of its base models.
                    'pred_time_test_marginal': The inference time of the model for the dataset provided, minus the inference time for the model's base models, if it has any.
                        Note that this ignores the time required to load the model into memory when bagging is disabled.
                If str is passed, `dataset` will be loaded using the str value as the file path.
            only_pareto_frontier : bool (optional)
                If `True`, only return model information of models in the Pareto frontier of the accuracy/latency trade-off (models which achieve the highest score within their end-to-end inference time).
                At minimum this will include the model with the highest score and the model with the lowest inference time.
                This is useful when deciding which model to use during inference if inference time is a consideration.
                Models filtered out by this process would never be optimal choices for a user that only cares about model inference time and score.
            silent : bool (optional)
                Should leaderboard DataFrame be printed?

            Returns
            -------
            Pandas `pandas.DataFrame` of model performance summary information.
        """
        dataset = self.__get_dataset(dataset) if dataset is not None else dataset
        return self._learner.leaderboard(X=dataset, only_pareto_frontier=only_pareto_frontier, silent=silent)

    def fit_summary(self, verbosity=3):
        """
            Output summary of information about models produced during `fit()`.
            May create various generated summary plots and store them in folder: `Predictor.output_directory`.

            Parameters
            ----------
            verbosity : int, default = 3
                Controls how detailed of a summary to ouput.
                Set <= 0 for no output printing, 1 to print just high-level summary,
                2 to print summary and create plots, >= 3 to print all information produced during fit().

            Returns
            -------
            Dict containing various detailed information. We do not recommend directly printing this dict as it may be very large.
        """
        hpo_used = len(self._trainer.hpo_results) > 0
        model_typenames = {key: self._trainer.model_types[key].__name__ for key in self._trainer.model_types}
        model_innertypenames = {key: self._trainer.model_types_inner[key].__name__ for key in self._trainer.model_types if key in self._trainer.model_types_inner}
        MODEL_STR = 'Model'
        ENSEMBLE_STR = 'Ensemble'
        for model in model_typenames:
            if (model in model_innertypenames) and (ENSEMBLE_STR not in model_innertypenames[model]) and (ENSEMBLE_STR in model_typenames[model]):
                new_model_typename = model_typenames[model] + "_" + model_innertypenames[model]
                if new_model_typename.endswith(MODEL_STR):
                    new_model_typename = new_model_typename[:-len(MODEL_STR)]
                model_typenames[model] = new_model_typename

        unique_model_types = set(model_typenames.values())  # no more class info
        # all fit() information that is returned:
        results = {
            'model_types': model_typenames,  # dict with key = model-name, value = type of model (class-name)
            'model_performance': self.model_performance,  # dict with key = model-name, value = validation performance
            'model_best': self._trainer.model_best,  # the name of the best model (on validation data)
            'model_paths': self._trainer.model_paths,  # dict with key = model-name, value = path to model file
            'model_fit_times': self._trainer.model_fit_times,
            'model_pred_times': self._trainer.model_pred_times,
            'num_bagging_folds': self._trainer.kfolds,
            'stack_ensemble_levels': self._trainer.stack_ensemble_levels,
            'feature_prune': self._trainer.feature_prune,
            'hyperparameter_tune': hpo_used,
            'hyperparameters_userspecified': self._trainer.hyperparameters,
        }
        if self.problem_type != REGRESSION:
            results['num_classes'] = self._trainer.num_classes
        if hpo_used:
            results['hpo_results'] = self._trainer.hpo_results
        # get dict mapping model name to final hyperparameter values for each model:
        model_hyperparams = {}
        for model_name in results['model_performance']:
            model_obj = self._trainer.load_model(model_name)
            model_hyperparams[model_name] = model_obj.params
        results['model_hyperparams'] = model_hyperparams

        if verbosity > 0:  # print stuff
            print("*** Summary of fit() ***")
            print("Estimated performance of each model:")
            results['leaderboard'] = self._learner.leaderboard(silent=False)
            # self._summarize('model_performance', 'Validation performance of individual models', results)
            #  self._summarize('model_best', 'Best model (based on validation performance)', results)
            # self._summarize('hyperparameter_tune', 'Hyperparameter-tuning used', results)
            print("Number of models trained: %s" % len(results['model_performance']))
            print("Types of models trained:")
            print(unique_model_types)
            num_fold_str = ""
            bagging_used = results['num_bagging_folds'] > 0
            if bagging_used:
                num_fold_str = f" (with {results['num_bagging_folds']} folds)"
            print("Bagging used: %s %s" % (bagging_used, num_fold_str))
            num_stack_str = ""
            stacking_used = results['stack_ensemble_levels'] > 0
            if stacking_used:
                num_stack_str = f" (with {results['stack_ensemble_levels']} levels)"
            print("Stack-ensembling used: %s %s" % (stacking_used, num_stack_str))
            hpo_str = ""
            if hpo_used and verbosity <= 2:
                hpo_str = " (call fit_summary() with verbosity >= 3 to see detailed HPO info)"
            print("Hyperparameter-tuning used: %s %s" % (hpo_used, hpo_str))
            # TODO: uncomment once feature_prune is functional:  self._summarize('feature_prune', 'feature-selection used', results)
            print("User-specified hyperparameters:")
            print(results['hyperparameters_userspecified'])
        if verbosity > 1:  # create plots
            plot_tabular_models(results, output_directory=self.output_directory,
                                save_file="SummaryOfModels.html",
                                plot_title="Models produced during fit()")
            if hpo_used:
                for model_type in results['hpo_results']:
                    plot_summary_of_models(
                        results['hpo_results'][model_type],
                        output_directory=self.output_directory, save_file=model_type + "_HPOmodelsummary.html",
                        plot_title=f"Models produced during {model_type} HPO")
                    plot_performance_vs_trials(
                        results['hpo_results'][model_type],
                        output_directory=self.output_directory, save_file=model_type + "_HPOperformanceVStrials.png",
                        plot_title=f"HPO trials for {model_type} models")
        if verbosity > 2:  # print detailed information
            if hpo_used:
                hpo_results = results['hpo_results']
                print("*** Details of Hyperparameter optimization ***")
                for model_type in hpo_results:
                    hpo_model = hpo_results[model_type]
                    print("HPO for %s model:  Num. configurations tried = %s, Time spent = %s, Search strategy = %s"
                          % (model_type, len(hpo_model['trial_info']), hpo_model['total_time'], hpo_model['search_strategy']))
                    print("Best hyperparameter-configuration (validation-performance: %s = %s):"
                          % (self.eval_metric, hpo_model['validation_performance']))
                    print(hpo_model['best_config'])
            """
            if bagging_used:
                pass # TODO: print detailed bagging info
            if stacking_used:
                pass # TODO: print detailed stacking info, like how much it improves validation performance
            if results['feature_prune']:
                pass # TODO: print detailed feature-selection info once feature-selection is functional.
            """
        if verbosity > 0:
            print("*** End of fit() summary ***")
        return results

    def transform_features(self, dataset=None, model=None, base_models=None):
        """
        Transforms dataset features through the AutoGluon feature generator.
        This is useful to gain an understanding of how AutoGluon interprets the dataset features.
        The output of this function can be used to train further models, even outside of AutoGluon.
        This can be useful for training your own models on the same data representation as AutoGluon.
        Individual AutoGluon models like the neural network may apply additional feature transformations that are not reflected in this method.
        This method only applies universal transforms employed by all AutoGluon models.
        When `dataset=None`, `base_models=[{best_model}], and bagging was enabled during fit():
            This returns the out-of-fold predictions of the best model, which can be used as training input to a custom user stacker model.

        Parameters
        ----------
        dataset : str or :class:`TabularDataset` or `pandas.DataFrame` (optional)
            The dataset to apply feature transformation to.
            This dataset does not require the label column.
            If str is passed, `dataset` will be loaded using the str value as the file path.
            If not specified, the original dataset used during fit() will be used if fit() was previously called with `cache_data=True`. Otherwise, an exception will be raised.
                For non-bagged mode predictors:
                    The dataset used when not specified is the validation set.
                    This can either be an automatically generated validation set or the user-defined `tuning_data` if passed during fit().
                    If all parameters are unspecified, then the output is equivalent to `predictor.load_data_internal(dataset='val', return_X=True, return_y=False)[0]`.
                    To get the label values of the output, call `predictor.load_data_internal(dataset='val', return_X=False, return_y=True)[1]`.
                    If the original training set is desired, it can be passed in through `dataset`.
                        Warning: Do not pass the original training set if `model` or `base_models` are set. This will result in overfit feature transformation.
                For bagged mode predictors:
                    The dataset used when not specified is the full training set.
                    If all parameters are unspecified, then the output is equivalent to `predictor.load_data_internal(dataset='train', return_X=True, return_y=False)[0]`.
                    To get the label values of the output, call `predictor.load_data_internal(dataset='train', return_X=False, return_y=True)[1]`.
                    `base_model` features generated in this instance will be from out-of-fold predictions.
                    Note that the training set may differ from the training set originally passed during fit(), as AutoGluon may choose to drop or duplicate rows during training.
                    Warning: Do not pass the original training set through `dataset` if `model` or `base_models` are set. This will result in overfit feature transformation. Instead set `dataset=None`.
        model : str, default = None
            Model to generate input features for.
            The output data will be equivalent to the input data that would be sent into `model.predict_proba(data)`.
                Note: This only applies to cases where `dataset` is not the training data.
            If `None`, then only return generically preprocessed features prior to any model fitting.
            Valid models are listed in this `predictor` by calling `predictor.model_names`.
            Specifying a `refit_full` model will cause an exception if `dataset=None`.
            `base_models=None` is a requirement when specifying `model`.
        base_models : list, default = None
            List of model names to use as base_models for a hypothetical stacker model when generating input features.
            If `None`, then only return generically preprocessed features prior to any model fitting.
            Valid models are listed in this `predictor` by calling `predictor.model_names`.
            If a stacker model S exists with `base_models=M`, then setting `base_models=M` is equivalent to setting `model=S`.
            `model=None` is a requirement when specifying `base_models`.

        Returns
        -------
        Pandas `pandas.DataFrame` of the provided `dataset` after feature transformation has been applied.
        This output does not include the label column, and will remove it if present in the supplied `dataset`.
        If a transformed label column is desired, use `predictor.transform_labels`.

        Examples
        --------
        >>> from autogluon import TabularPrediction as task
        >>> train_data = task.Dataset('train.csv')
        >>> predictor = task.fit(train_data=train_data, label='class', auto_stack=True, cache_data=True)  # predictor is in bagged mode and `cache_data=True`.
        >>> model = 'weighted_ensemble_k0_l1'
        >>> test_data = task.Dataset('test.csv')
        >>> train_data_transformed = predictor.transform_features(model=model)  # Internal training DataFrame used as input to `model.fit()` during `predictor = task.fit(train_data=train_data, ...)`
        >>> test_data_transformed = predictor.transform_features(dataset=test_data, model=model)  # Internal test DataFrame used as input to `model.predict_proba()` during `predictor.predict_proba(test_data, model=model)`

        """
        dataset = self.__get_dataset(dataset) if dataset is not None else dataset
        # TODO: Make this index fix inside of learner/trainer once the defect is identified. For now, this resolves the issue.
        if dataset is not None:
            original_indices = copy.deepcopy(dataset.index)
            dataset = dataset.reset_index(drop=True)
        else:
            original_indices = None

        dataset_transformed = self._learner.get_inputs_to_stacker(dataset=dataset, model=model, base_models=base_models)
        if original_indices is not None:
            dataset_transformed.index = original_indices
        return dataset_transformed

    # TODO: Add support for DataFrame input and support for DataFrame output
    #  Add support for retaining DataFrame/Series indices in output
    #  Add as_pandas parameter
    def transform_labels(self, labels, inverse=False, proba=False):
        """
        Transforms dataset labels to the internal label representation.
        This can be useful for training your own models on the same data label representation as AutoGluon.
        Regression problems do not differ between original and internal representation, and thus this method will return the provided labels.
        Warning: When `inverse=False`, it is possible for the output to contain NaN label values in multiclass problems if the provided label was dropped during training.

        Parameters
        ----------
        labels : `numpy.ndarray` or `pandas.Series`
            Labels to transform.
            If `proba=False`, an example input would be the output of `predictor.predict(test_data)`.
            If `proba=True`, an example input would be the output of `predictor.predict_proba(test_data)`.
        inverse : boolean, default = False
            When `True`, the input labels are treated as being in the internal representation and the original representation is outputted.
        proba : boolean, default = False
            When `True`, the input labels are treated as probabilities and the output will be the internal representation of probabilities.
                In this case, it is expected that `labels` be a `numpy.ndarray`.
                If the `problem_type` is multiclass:
                    The input column order must be equal to `predictor.class_labels`.
                    The output column order will be equal to `predictor.class_labels_internal`.
                    if `inverse=True`, the same logic applies, but with input and output columns interchanged.
            When `False`, the input labels are treated as actual labels and the output will be the internal representation of the labels.
                In this case, it is expected that `labels` be a `pandas.Series` or `numpy.ndarray`.

        Returns
        -------
        Pandas `pandas.Series` of labels if `proba=False` or Numpy `numpy.ndarray` of label probabilities if `proba=True`

        """
        if inverse:
            if proba:
                labels_transformed = self._learner.label_cleaner.inverse_transform_proba(y=labels)
            else:
                labels_transformed = self._learner.label_cleaner.inverse_transform(y=labels)
        else:
            if proba:
                labels_transformed = self._learner.label_cleaner.transform_proba(y=labels)
            else:
                labels_transformed = self._learner.label_cleaner.transform(y=labels)
        return labels_transformed

    # TODO: Consider adding time_limit option to early stop the feature importance process
    def feature_importance(self, model=None, dataset=None, features=None, raw=True, subsample_size=10000, silent=False):
        """
        Calculates feature importance scores for the given model.
        A feature's importance score represents the performance drop that results when the model makes predictions on a perturbed copy of the dataset where this feature's values have been randomly shuffled across rows.
        A feature score of 0.01 would indicate that the predictive performance dropped by 0.01 when the feature was randomly shuffled.
        The higher the score a feature has, the more important it is to the model's performance.
        If a feature has a negative score, this means that the feature is likely harmful to the final model, and a model trained with the feature removed would be expected to achieve a better predictive performance.
        Note that calculating feature importance can be a very computationally expensive process, particularly if the model uses hundreds or thousands of features. In many cases, this can take longer than the original model training.
        To estimate how long `feature_importance(model, dataset, features)` will take, it is roughly the time taken by `predict_proba(dataset, model)` multiplied by the number of features.

        Parameters
        ----------
        model : str, default = None
            Model to get feature importances for, if None the best model is chosen.
            Valid models are listed in this `predictor` by calling `predictor.model_names`
        dataset : str or :class:`TabularDataset` or `pandas.DataFrame` (optional)
            This dataset must also contain the label-column with the same column-name as specified during fit().
            If specified, then the dataset is used to calculate the feature importance scores.
            If str is passed, `dataset` will be loaded using the str value as the file path.
            If not specified, the original dataset used during fit() will be used if `cache_data=True`. Otherwise, an exception will be raised.
            Do not pass the training data through this argument, as the feature importance scores calculated will be inaccurate.
        features : list, default = None
            List of str feature names that feature importances are calculated for and returned, specify None to get all feature importances.
            If you only want to compute feature importances for some of the features, you can pass their names in as a list of str.
        raw : bool, default = True
            Whether to compute feature importance on raw features in the original data (after automated feature engineering) or on the features used by the particular model.
            For example, a stacker model uses both the original features and the predictions of the lower-level models.
            Note that for bagged models, feature importance calculation is not yet supported when both `raw=True` and `dataset=None`. Doing so will raise an exception.
        subsample_size : int, default = 10000
            The number of rows to sample from `dataset` when computing feature importance.
            If `subsample_size=None` or `dataset` contains fewer than `subsample_size` rows, all rows will be used during computation.
            Larger values increase the accuracy of the feature importance scores.
            Runtime linearly scales with `subsample_size`.
        silent : bool, default = False
            Whether to suppress logging output

        Returns
        -------
        Pandas `pandas.Series` of feature importance scores.

        """
        dataset = self.__get_dataset(dataset) if dataset is not None else dataset
        if (dataset is None) and (not self._trainer.is_data_saved):
            raise AssertionError('No dataset was provided and there is no cached data to load for feature importance calculation. `cache_data=True` must be set in the `TabularPrediction.fit()` call to enable this functionality when dataset is not specified.')

        return self._learner.get_feature_importance(model=model, X=dataset, features=features, raw=raw, subsample_size=subsample_size, silent=silent)

    def refit_full(self, model='all'):
        """
        Retrain model on all of the data (training + validation).
        For bagged models:
            Optimizes a model's inference time by collapsing bagged ensembles into a single model fit on all of the training data.
            This process will typically result in a slight accuracy reduction and a large inference speedup.
            The inference speedup will generally be between 10-200x faster than the original bagged ensemble model.
                The inference speedup factor is equivalent to (k * n), where k is the number of folds (`num_bagging_folds`) and n is the number of finished repeats (`num_bagging_sets`) in the bagged ensemble.
            The runtime is generally 10% or less of the original fit runtime.
                The runtime can be roughly estimated as 1 / (k * n) of the original fit runtime, with k and n defined above.
        For non-bagged models:
            Optimizes a model's accuracy by retraining on 100% of the data without using a validation set.
            Will typically result in a slight accuracy increase and no change to inference time.
            The runtime will be approximately equal to the original fit runtime.
        This process does not alter the original models, but instead adds additional models.
        If stacker models are refit by this process, they will use the refit_full versions of the ancestor models during inference.
        Models produced by this process will not have validation scores, as they use all of the data for training.
            Therefore, it is up to the user to determine if the models are of sufficient quality by including test data in `predictor.leaderboard(dataset=test_data)`.
            If the user does not have additional test data, they should reference the original model's score for an estimate of the performance of the refit_full model.
                Warning: Be aware that utilizing refit_full models without separately verifying on test data means that the model is untested, and has no guarantee of being consistent with the original model.
        `cache_data` must have been set to `True` during the original training to enable this functionality.

        Parameters
        ----------
        model : str, default = 'all'
            Model name of model to refit.
                If 'all' then all models are refitted.
                If 'best' then the model with the highest validation score is refit.
            All ancestor models will also be refit in the case that the selected model is a weighted or stacker ensemble.
            Valid models are listed in this `predictor` by calling `predictor.model_names`.

        Returns
        -------
        Dictionary of original model names -> refit_full model names.
        """
        refit_full_dict = self._learner.refit_ensemble_full(model=model)
        self.model_names = self._trainer.get_model_names_all()
        return refit_full_dict

    def get_model_full_dict(self):
        """
        Returns a dictionary of original model name -> refit full model name.
        Empty unless `refit_full=True` was set during fit or `predictor.refit_full()` was called.
        This can be useful when determining the best model based off of `predictor.leaderboard()`, then getting the _FULL version of the model by passing its name as the key to this dictionary.

        Returns
        -------
        Dictionary of original model name -> refit full model name.
        """
        return copy.deepcopy(self._trainer.model_full_dict)

    def info(self):
        """
        [EXPERIMENTAL] Returns a dictionary of `predictor` metadata.
        Warning: This functionality is currently in preview mode.
            The metadata information returned may change in structure in future versions without warning.
            The definitions of various metadata values are not yet documented.
            The output of this function should not be used for programmatic decisions.
        Contains information such as row count, column count, model training time, validation scores, hyperparameters, and much more.

        Returns
        -------
        Dictionary of `predictor` metadata.
        """
        return self._learner.get_info(include_model_info=True)

    @classmethod
    def load(cls, output_directory, verbosity=2):
        """
        Load a predictor object previously produced by `fit()` from file and returns this object.
        Is functionally equivalent to :meth:`autogluon.task.tabular_prediction.TabularPrediction.load`.

        Parameters
        ----------
        output_directory : str
            Path to directory where trained models are stored (i.e. the `output_directory` specified in previous call to `fit()`).
        verbosity : int, default = 2
            Verbosity levels range from 0 to 4 and control how much information is generally printed by this Predictor.
            Higher levels correspond to more detailed print statements (you can set verbosity = 0 to suppress warnings).
            If using logging, you can alternatively control amount of information printed via `logger.setLevel(L)`,
            where `L` ranges from 0 to 50 (Note: higher values `L` correspond to fewer print statements, opposite of verbosity levels)

        Returns
        -------
        :class:`TabularPredictor` object
        """
        logger.setLevel(verbosity2loglevel(verbosity))  # Reset logging after load (may be in new Python session)
        if output_directory is None:
            raise ValueError("output_directory cannot be None in load()")

        output_directory = setup_outputdir(output_directory)  # replace ~ with absolute path if it exists
        learner = Learner.load(output_directory)
        return cls(learner=learner)

    def save(self):
        """ Save this predictor to file in directory specified by this Predictor's `output_directory`.
            Note that `fit()` already saves the predictor object automatically
            (we do not recommend modifying the Predictor object yourself as it tracks many trained models).
        """
        self._learner.save()
        logger.log(20, "TabularPredictor saved. To load, use: TabularPredictor.load(\"%s\")" % self.output_directory)

    def load_data_internal(self, dataset='train', return_X=True, return_y=True):
        """
        Loads the internal data representation used during model training.
        Individual AutoGluon models like the neural network may apply additional feature transformations that are not reflected in this method.
        This method only applies universal transforms employed by all AutoGluon models.
        This will raise an exception if `cache_data=False` was set in `task.fit()`.
        Warning, the internal representation may:
            Have different features compared to the original data.
            Have different row counts compared to the original data.
            Have indices which do not align with the original data.
            Have label values which differ from those in the original data.
        Internal data representations should NOT be combined with the original data, in most cases this is not possible.

        Parameters
        ----------
        dataset : str, default = 'train'
            The dataset to load.
            Valid values are:
                'train':
                    Load the training data used during model training.
                    This is a transformed and augmented version of the `train_data` passed in `task.fit()`.
                'val':
                    Load the validation data used during model training.
                    This is a transformed and augmented version of the `tuning_data` passed in `task.fit()`.
                    If `tuning_data=None` was set in `task.fit()`, then `tuning_data` is an automatically generated validation set created by splitting `train_data`.
                    Warning: Will raise an exception if called by a bagged predictor, as bagged predictors have no validation data.
        return_X : bool, default = True
            Whether to return the internal data features
            If set to `False`, then the first element in the returned tuple will be None.
        return_y : bool, default = True
            Whether to return the internal data labels
            If set to `False`, then the second element in the returned tuple will be None.

        Returns
        -------
        Tuple of (`pandas.DataFrame`, `pandas.Series`) corresponding to the internal data features and internal data labels, respectively.

        """
        if dataset == 'train':
            load_X = self._trainer.load_X_train
            load_y = self._trainer.load_y_train
        elif dataset == 'val':
            load_X = self._trainer.load_X_val
            load_y = self._trainer.load_y_val
        else:
            raise ValueError(f'dataset must be one of: [\'train\', \'val\'], but was \'{dataset}\'.')
        X = load_X() if return_X else None
        y = load_y() if return_y else None
        return X, y

    @staticmethod
    def _summarize(key, msg, results):
        if key in results:
            print(msg + ": " + str(results[key]))

    @staticmethod
    def __get_dataset(dataset):
        if isinstance(dataset, TabularDataset):
            return dataset
        elif isinstance(dataset, pd.DataFrame):
            return TabularDataset(df=dataset)
        elif isinstance(dataset, str):
            return TabularDataset(file_path=dataset)
        elif isinstance(dataset, pd.Series):
            raise TypeError("dataset must be TabularDataset or pandas.DataFrame, not pandas.Series. \
                   To predict on just single example (ith row of table), use dataset.iloc[[i]] rather than dataset.iloc[i]")
        else:
            raise TypeError("dataset must be TabularDataset or pandas.DataFrame or str file path to dataset")
