import datetime
import importlib
import logging
import os
import re
import shutil
import sys
import warnings

import cloudpickle
import yaml

import mlflow
from mlflow.entities import SourceType, ViewType
from mlflow.environment_variables import MLFLOW_RECIPES_EXECUTION_TARGET_STEP_NAME
from mlflow.exceptions import BAD_REQUEST, INVALID_PARAMETER_VALUE, MlflowException
from mlflow.models import Model
from mlflow.recipes.artifacts import (
    DataframeArtifact,
    HyperParametersArtifact,
    ModelArtifact,
    RunArtifact,
)
from mlflow.recipes.cards import BaseCard
from mlflow.recipes.step import BaseStep, StepClass
from mlflow.recipes.utils.execution import get_step_output_path
from mlflow.recipes.utils.metrics import (
    _get_builtin_metrics,
    _get_custom_metrics,
    _get_error_fn,
    _get_extended_task,
    _get_model_type_from_template,
    _get_primary_metric,
    _load_custom_metrics,
    transform_multiclass_metrics_dict,
)
from mlflow.recipes.utils.step import (
    get_merged_eval_metrics,
    get_pandas_data_profiles,
    validate_classification_config,
)
from mlflow.recipes.utils.tracking import (
    TrackingConfig,
    apply_recipe_tracking_config,
    get_recipe_tracking_config,
    get_run_tags_env_vars,
    log_code_snapshot,
)
from mlflow.recipes.utils.wrapped_recipe_model import WrappedRecipeModel
from mlflow.tracking import MlflowClient
from mlflow.tracking.fluent import _get_experiment_id
from mlflow.utils.databricks_utils import get_databricks_env_vars, get_databricks_run_url
from mlflow.utils.file_utils import TempDir
from mlflow.utils.mlflow_tags import (
    MLFLOW_RECIPE_PROFILE_NAME,
    MLFLOW_RECIPE_STEP_NAME,
    MLFLOW_RECIPE_TEMPLATE_NAME,
    MLFLOW_SOURCE_TYPE,
)
from mlflow.utils.string_utils import strip_prefix

_REBALANCING_CUTOFF = 5000
_REBALANCING_DEFAULT_RATIO = 0.3
_USER_DEFINED_TRAIN_STEP_MODULE = "steps.train"

_logger = logging.getLogger(__name__)


class TrainStep(BaseStep):
    MODEL_ARTIFACT_RELATIVE_PATH = "model"
    SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH = "sk_model"
    PREDICTED_TRAINING_DATA_RELATIVE_PATH = "predicted_training_data.parquet"

    def __init__(self, step_config, recipe_root, recipe_config=None):
        super().__init__(step_config, recipe_root)
        self.tracking_config = TrackingConfig.from_dict(self.step_config)
        self.recipe_config = recipe_config

    def _validate_and_apply_step_config(self):
        if "using" in self.step_config:
            if self.step_config["using"] not in ["custom", "automl/flaml"]:
                raise MlflowException(
                    f"Invalid train step configuration value {self.step_config['using']} for "
                    f"key 'using'. Supported values are: ['custom', 'automl/flaml']",
                    error_code=INVALID_PARAMETER_VALUE,
                )
        else:
            self.step_config["using"] = "custom"

        if "tuning" in self.step_config:
            if "enabled" in self.step_config["tuning"] and isinstance(
                self.step_config["tuning"]["enabled"], bool
            ):
                self.step_config["tuning_enabled"] = self.step_config["tuning"]["enabled"]
            else:
                raise MlflowException(
                    "The 'tuning' configuration in the train step must include an "
                    "'enabled' key whose value is either true or false.",
                    error_code=INVALID_PARAMETER_VALUE,
                )

            if self.step_config["tuning_enabled"]:
                if "sample_fraction" in self.step_config["tuning"]:
                    sample_fraction = float(self.step_config["tuning"]["sample_fraction"])
                    if 0 < sample_fraction <= 1.0:
                        self.step_config["sample_fraction"] = sample_fraction
                    else:
                        raise MlflowException(
                            "The tuning 'sample_fraction' configuration in the train step "
                            "must be between 0 and 1.",
                            error_code=INVALID_PARAMETER_VALUE,
                        )
                else:
                    self.step_config["sample_fraction"] = 1.0

                if "algorithm" not in self.step_config["tuning"]:
                    self.step_config["tuning"]["algorithm"] = "hyperopt.rand.suggest"

                if "parallelism" not in self.step_config["tuning"]:
                    self.step_config["tuning"]["parallelism"] = 1

                if "max_trials" not in self.step_config["tuning"]:
                    raise MlflowException(
                        "The 'max_trials' configuration in the train step must be provided.",
                        error_code=INVALID_PARAMETER_VALUE,
                    )

                if "parameters" not in self.step_config["tuning"]:
                    raise MlflowException(
                        "The 'parameters' configuration in the train step must be provided "
                        " when tuning is enabled.",
                        error_code=INVALID_PARAMETER_VALUE,
                    )

        else:
            self.step_config["tuning_enabled"] = False

        if "estimator_params" not in self.step_config:
            self.step_config["estimator_params"] = {}

        self.target_col = self.step_config.get("target_col")
        if self.target_col is None:
            raise MlflowException(
                "Missing target_col config in recipe config.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        self.recipe = self.step_config.get("recipe")
        if self.recipe is None:
            raise MlflowException(
                "Missing recipe config in recipe config.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        self.positive_class = self.step_config.get("positive_class")
        self.extended_task = _get_extended_task(self.recipe, self.positive_class)
        self.rebalance_training_data = self.step_config.get("rebalance_training_data", True)
        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
        if "estimator_method" not in self.step_config and self.step_config["using"] == "custom":
            raise MlflowException(
                "Missing 'estimator_method' configuration in the train step, "
                "which is using 'custom'.",
                error_code=INVALID_PARAMETER_VALUE,
            )

        self.predict_prefix = self.step_config.get("predict_prefix", "predicted_")
        self.predict_scores_for_all_classes = self.step_config.get(
            "predict_scores_for_all_classes", True
        )
        self.primary_metric = _get_primary_metric(
            self.step_config.get("primary_metric"), self.extended_task
        )
        builtin_metrics = _get_builtin_metrics(self.extended_task)
        custom_metrics = _get_custom_metrics(self.step_config, self.extended_task)
        self.user_defined_custom_metrics = {metric.name: metric for metric in custom_metrics}
        self.evaluation_metrics = {metric.name: metric for metric in builtin_metrics}
        self.evaluation_metrics.update(self.user_defined_custom_metrics)
        self.evaluation_metrics_greater_is_better = {
            metric.name: metric.greater_is_better for metric in builtin_metrics
        }
        self.evaluation_metrics_greater_is_better.update(
            {metric.name: metric.greater_is_better for metric in custom_metrics}
        )
        if self.primary_metric is not None and self.primary_metric not in self.evaluation_metrics:
            raise MlflowException(
                f"The primary metric '{self.primary_metric}' is a custom metric, but its"
                " corresponding custom metric configuration is missing from `recipe.yaml`.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        self.code_paths = [os.path.join(self.recipe_root, "steps")]

    @classmethod
    def construct_search_space_from_yaml(cls, params):
        from hyperopt import hp

        search_space = {}
        for param_name, param_details in params.items():
            if "values" in param_details:
                param_details_to_pass = param_details["values"]
                search_space[param_name] = hp.choice(param_name, param_details_to_pass)
            elif "distribution" in param_details:
                hp_tuning_fn = getattr(hp, param_details["distribution"])
                param_details_to_pass = param_details.copy()
                param_details_to_pass.pop("distribution")
                search_space[param_name] = hp_tuning_fn(param_name, **param_details_to_pass)
            else:
                raise MlflowException(
                    f"Parameter {param_name} must contain either a list of 'values' or a "
                    f"'distribution' following hyperopt parameter expressions",
                    error_code=INVALID_PARAMETER_VALUE,
                )
        return search_space

    @classmethod
    def is_tuning_param_equal(cls, tuning_param, logged_param):
        if isinstance(tuning_param, bool):
            return tuning_param == bool(logged_param)
        elif isinstance(tuning_param, int):
            return tuning_param == int(logged_param)
        elif isinstance(tuning_param, float):
            return tuning_param == float(logged_param)
        elif isinstance(tuning_param, str):
            return tuning_param.strip() == logged_param.strip()
        else:
            return tuning_param == logged_param

    def _fitted_estimator(self, *args):
        if "calibrate_proba" in self.step_config:
            return self._calibrated_classifier_fitted_estimator(*args)
        else:
            return self._label_encoded_fitted_estimator(*args)

    def _calibrated_classifier_fitted_estimator(self, estimator, X_train, y_train):
        original_estimator = estimator.fit(X_train, y_train)
        if "classification" in self.recipe:
            from sklearn.calibration import CalibratedClassifierCV

            estimator = CalibratedClassifierCV(
                estimator, method=self.step_config["calibrate_proba"]
            )
            estimator.fit(X_train, y_train)

        return estimator, {"original_estimator": original_estimator}

    def _label_encoded_fitted_estimator(self, estimator, X_train, y_train):
        from sklearn.preprocessing import LabelEncoder

        label_encoder = None
        target_column_class_labels = None
        encoded_y_train = y_train
        if "classification" in self.recipe:
            label_encoder = LabelEncoder()
            label_encoder.fit(y_train)
            target_column_class_labels = label_encoder.classes_
            import pandas as pd

            encoded_y_train = pd.Series(label_encoder.transform(y_train))

        def inverse_label_encoder(predicted_output):
            if not label_encoder:
                return predicted_output

            return label_encoder.inverse_transform(predicted_output)

        estimator.fit(X_train, encoded_y_train)
        original_predict = estimator.predict

        def wrapped_predict(*args, **kwargs):
            return inverse_label_encoder(original_predict(*args, **kwargs))

        estimator.predict = wrapped_predict

        return estimator, {"target_column_class_labels": target_column_class_labels}

    def _run(self, output_directory):
        def my_warn(*args, **kwargs):
            timestamp = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
            stacklevel = 1 if "stacklevel" not in kwargs else kwargs["stacklevel"]
            frame = sys._getframe(stacklevel)
            filename = frame.f_code.co_filename
            lineno = frame.f_lineno
            message = f"{timestamp} {filename}:{lineno}: {args[0]}\n"
            with open(os.path.join(output_directory, "warning_logs.txt"), "a") as f:
                f.write(message)

        original_warn = warnings.warn
        warnings.warn = my_warn
        try:
            import numpy as np
            import pandas as pd
            import sklearn
            from sklearn.pipeline import make_pipeline
            from sklearn.utils.class_weight import compute_class_weight

            from mlflow.models import infer_signature

            with open(os.path.join(output_directory, "warning_logs.txt"), "w"):
                pass

            apply_recipe_tracking_config(self.tracking_config)

            transformed_training_data_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="transform",
                relative_path="transformed_training_data.parquet",
            )
            train_df = pd.read_parquet(transformed_training_data_path)
            validate_classification_config(
                self.task, self.positive_class, train_df, self.target_col
            )
            self.using_rebalancing = False
            if self.extended_task == "classification/binary":
                classes = np.unique(train_df[self.target_col])
                class_weights = compute_class_weight(
                    class_weight="balanced",
                    classes=classes,
                    y=train_df[self.target_col],
                )
                self.original_class_weights = dict(zip(classes, class_weights))
                if self.rebalance_training_data and len(classes) == 2:
                    if len(train_df) > _REBALANCING_CUTOFF:
                        self.using_rebalancing = True
                        train_df = self._rebalance_classes(train_df)
                    else:
                        _logger.info(
                            f"Training data has less than {_REBALANCING_CUTOFF} rows, "
                            f"skipping rebalancing."
                        )

            X_train, y_train = train_df.drop(columns=[self.target_col]), train_df[self.target_col]

            transformed_validation_data_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="transform",
                relative_path="transformed_validation_data.parquet",
            )
            validation_df = pd.read_parquet(transformed_validation_data_path)

            raw_training_data_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="split",
                relative_path="train.parquet",
            )
            raw_train_df = pd.read_parquet(raw_training_data_path)
            raw_X_train = raw_train_df.drop(columns=[self.target_col])

            raw_validation_data_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="split",
                relative_path="validation.parquet",
            )
            raw_validation_df = pd.read_parquet(raw_validation_data_path)

            transformer_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="transform",
                relative_path="transformer.pkl",
            )

            tags = {
                MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.RECIPE),
                MLFLOW_RECIPE_TEMPLATE_NAME: self.step_config["recipe"],
                MLFLOW_RECIPE_PROFILE_NAME: self.step_config["profile"],
                MLFLOW_RECIPE_STEP_NAME: MLFLOW_RECIPES_EXECUTION_TARGET_STEP_NAME.get(),
            }

            run_name = self.tracking_config.run_name
            best_estimator_params = None
            mlflow.autolog(log_models=False, silent=True)
            with mlflow.start_run(run_name=run_name, tags=tags) as run:
                estimator = self._resolve_estimator(
                    X_train, y_train, validation_df, run, output_directory
                )
                best_estimator_params = estimator.get_params()
                fitted_estimator, additional_fitted_args = self._fitted_estimator(
                    estimator, X_train, y_train
                )
                logged_estimator = self._log_estimator_to_mlflow(fitted_estimator, X_train)

                # Create a recipe consisting of the transformer+model for test data evaluation
                with open(transformer_path, "rb") as f:
                    transformer = cloudpickle.load(f)
                mlflow.sklearn.log_model(
                    transformer, "transform/transformer", code_paths=self.code_paths
                )

                trained_pipeline = make_pipeline(transformer, fitted_estimator)
                # Creating a wrapped recipe model which exposes a single predict function
                # so it can output both predict and predict_proba(for a classification problem)
                # at the same time.
                wrapped_model = WrappedRecipeModel(
                    self.predict_scores_for_all_classes,
                    self.predict_prefix,
                    target_column_class_labels=additional_fitted_args.get(
                        "target_column_class_labels"
                    ),
                )

                model_uri = get_step_output_path(
                    recipe_root_path=self.recipe_root,
                    step_name=self.name,
                    relative_path=TrainStep.MODEL_ARTIFACT_RELATIVE_PATH,
                )
                sklearn_model_uri = get_step_output_path(
                    recipe_root_path=self.recipe_root,
                    step_name=self.name,
                    relative_path=TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH,
                )
                if os.path.exists(model_uri):
                    shutil.rmtree(model_uri)
                if os.path.exists(sklearn_model_uri):
                    shutil.rmtree(sklearn_model_uri)

                # Saving the sklearn model as a separate output since `mlflow.evaluate()`, which is
                # used in evaluate step of the recipe, needs this model's sklearn representation
                # to computes metrics (the pyfunc representation of the user-facing model logged to
                # MLflow Tracking is not currently compatible with `mlflow.evaluate()`)
                mlflow.sklearn.save_model(trained_pipeline, sklearn_model_uri)
                artifacts = {"model_path": sklearn_model_uri}
                with TempDir() as tmp:
                    # Saving a temp model so that the output schema (signature) of the model's
                    # pyfunc representation can be inferred and included when logging the model
                    # to MLflow Tracking. Unfortunately, there is currently no easy way to infer
                    # the model's signature without first saving a copy of it, and there is no easy
                    # way to add an inferred signature to an existing model
                    pyfunc_model_tmp_path = os.path.join(tmp.path(), "pyfunc_model")
                    mlflow.pyfunc.save_model(
                        path=pyfunc_model_tmp_path,
                        python_model=wrapped_model,
                        artifacts=artifacts,
                    )
                    tempModel = mlflow.pyfunc.load_model(pyfunc_model_tmp_path)
                    model_schema = infer_signature(
                        raw_X_train, tempModel.predict(raw_X_train.copy())
                    )
                    mlflow.pyfunc.save_model(
                        path=model_uri,
                        python_model=wrapped_model,
                        artifacts=artifacts,
                        signature=model_schema,
                        code_path=self.code_paths,
                    )
                model = mlflow.pyfunc.load_model(model_uri)
                # Adding a sklearn flavor to the pyfunc model so models could be loaded easily
                # using mlflow.sklearn.load_model
                tmp_model_info = Model.load(model_uri)
                model_data_subpath = os.path.join(
                    "artifacts", TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH, "model.pkl"
                )
                model_info = Model(
                    artifact_path="train/model",
                    run_id=run.info.run_id,
                    utc_time_created=tmp_model_info.utc_time_created,
                    flavors=tmp_model_info.flavors,
                    signature=tmp_model_info.signature,  # ModelSignature
                    saved_input_example_info=tmp_model_info.saved_input_example_info,
                    model_uuid=tmp_model_info.model_uuid,
                    mlflow_version=tmp_model_info.mlflow_version,
                    metadata=tmp_model_info.metadata,
                )
                model_info.add_flavor(
                    mlflow.sklearn.FLAVOR_NAME,
                    pickled_model=model_data_subpath,
                    sklearn_version=sklearn.__version__,
                    serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE,
                    code="code",
                )
                model_info.save(f"{model_uri}/MLmodel")
                mlflow.log_artifacts(model_uri, "train/model")

                with open(os.path.join(output_directory, "run_id"), "w") as f:
                    f.write(run.info.run_id)
                log_code_snapshot(
                    self.recipe_root, run.info.run_id, recipe_config=self.recipe_config
                )

                eval_metrics = {}
                for dataset_name, (dataset, metric_prefix) in {
                    "training": (train_df, "training_"),
                    "validation": (validation_df, "val_"),
                }.items():
                    eval_config = {
                        "log_model_explainability": False,
                        "metric_prefix": metric_prefix,
                    }
                    if self.positive_class is not None:
                        eval_config["pos_label"] = self.positive_class
                    eval_result = mlflow.evaluate(
                        model=logged_estimator.model_uri,
                        data=dataset,
                        targets=self.target_col,
                        model_type=_get_model_type_from_template(self.recipe),
                        evaluators="default",
                        extra_metrics=_load_custom_metrics(
                            self.recipe_root,
                            self.evaluation_metrics.values(),
                        ),
                        evaluator_config=eval_config,
                    )
                    eval_result.save(os.path.join(output_directory, f"eval_{dataset_name}"))
                    eval_metrics[dataset_name] = {
                        strip_prefix(k, metric_prefix): v for k, v in eval_result.metrics.items()
                    }

            target_data = raw_validation_df[self.target_col]
            prediction_result = model.predict(raw_validation_df.drop(self.target_col, axis=1))

            use_probability_for_error_rate = False
            if isinstance(prediction_result, pd.DataFrame) and {
                f"{self.predict_prefix}label",
                f"{self.predict_prefix}score",
            }.issubset(prediction_result.columns):
                if self.positive_class:
                    prediction_result_for_error = prediction_result[
                        f"{self.predict_prefix}score_{self.positive_class}"
                    ].values
                    # use_probability_for_error_rate to true to compute error function
                    # based on positive class
                    use_probability_for_error_rate = True
                else:
                    prediction_result_for_error = prediction_result[
                        f"{self.predict_prefix}score"
                    ].values

                prediction_result = prediction_result[f"{self.predict_prefix}label"].values
            else:
                prediction_result_for_error = prediction_result
            error_fn = _get_error_fn(
                self.recipe,
                use_probability=use_probability_for_error_rate,
                positive_class=self.positive_class,
            )
            pred_and_error_df = pd.DataFrame(
                {
                    "target": target_data,
                    "prediction": prediction_result,
                    "error": error_fn(prediction_result_for_error, target_data.to_numpy()),
                }
            )
            calibrated_plot = None
            train_predictions = model.predict(raw_train_df.drop(self.target_col, axis=1))
            if isinstance(train_predictions, pd.DataFrame) and {
                f"{self.predict_prefix}label",
                f"{self.predict_prefix}score",
            }.issubset(train_predictions.columns):
                predicted_training_data = raw_train_df.assign(
                    predicted_data=train_predictions[f"{self.predict_prefix}label"],
                    predicted_score=train_predictions[f"{self.predict_prefix}score"].values,
                )
                if self.positive_class:
                    worst_examples_df = BaseStep._generate_worst_examples_dataframe(
                        raw_train_df,
                        train_predictions[f"{self.predict_prefix}label"].values,
                        error_fn(
                            train_predictions[
                                f"{self.predict_prefix}score_{self.positive_class}"
                            ].values,
                            raw_train_df[self.target_col].to_numpy(),
                        ),
                        self.target_col,
                    )

                    if "calibrate_proba" in self.step_config and hasattr(
                        additional_fitted_args.get("original_estimator"), "predict_proba"
                    ):
                        from sklearn.calibration import CalibrationDisplay

                        calibrated_plot = CalibrationDisplay.from_estimator(
                            additional_fitted_args.get("original_estimator"),
                            raw_train_df.drop(self.target_col, axis=1),
                            raw_train_df[self.target_col],
                            pos_label=self.positive_class,
                        )
                else:
                    # compute worst examples data_frame only if positive class exists
                    worst_examples_df = pd.DataFrame()

            else:
                predicted_training_data = raw_train_df.assign(predicted_data=train_predictions)
                worst_examples_df = BaseStep._generate_worst_examples_dataframe(
                    raw_train_df,
                    train_predictions,
                    error_fn(train_predictions, raw_train_df[self.target_col].to_numpy()),
                    self.target_col,
                )
            predicted_training_data.to_parquet(
                os.path.join(output_directory, TrainStep.PREDICTED_TRAINING_DATA_RELATIVE_PATH)
            )

            leaderboard_df = None
            try:
                leaderboard_df = self._get_leaderboard_df(run, eval_metrics)
            except Exception as e:
                _logger.warning(
                    "Failed to build model leaderboard due to unexpected failure: %s", e
                )
            tuning_df = None
            if self.step_config["tuning_enabled"]:
                try:
                    tuning_df = self._get_tuning_df(run, params=best_estimator_params.keys())
                except Exception as e:
                    _logger.warning(
                        "Failed to build tuning results table due to unexpected failure: %s", e
                    )

            card = self._build_step_card(
                eval_metrics=eval_metrics,
                pred_and_error_df=pred_and_error_df,
                model_schema=model_schema,
                run_id=run.info.run_id,
                model_uri=model_uri,
                worst_examples_df=worst_examples_df,
                train_df=raw_train_df,
                output_directory=output_directory,
                leaderboard_df=leaderboard_df,
                tuning_df=tuning_df,
                calibrated_plot=calibrated_plot,
            )
            card.save_as_html(output_directory)
            for step_name in ("ingest", "split", "transform", "train"):
                self._log_step_card(run.info.run_id, step_name)
            return card
        finally:
            warnings.warn = original_warn

    def _get_user_defined_estimator(self, X_train, y_train, validation_df, run, output_directory):
        sys.path.append(self.recipe_root)
        estimator_fn = getattr(
            importlib.import_module(_USER_DEFINED_TRAIN_STEP_MODULE),
            self.step_config["estimator_method"],
        )
        estimator_hardcoded_params = self.step_config["estimator_params"]

        # if using rebalancing pass in original class weights to preserve original distribution
        if self.using_rebalancing:
            estimator_hardcoded_params["class_weight"] = self.original_class_weights

        if self.step_config["tuning_enabled"]:
            estimator_hardcoded_params, best_hp_params = self._tune_and_get_best_estimator_params(
                run.info.run_id,
                estimator_hardcoded_params,
                estimator_fn,
                X_train,
                y_train,
                validation_df,
            )
            best_combined_params = dict(estimator_hardcoded_params, **best_hp_params)
            estimator = estimator_fn(best_combined_params)
            all_estimator_params = estimator.get_params()
            default_params_keys = all_estimator_params.keys() - best_combined_params.keys()
            default_params = {k: all_estimator_params[k] for k in default_params_keys}
            self._write_best_parameters_outputs(
                output_directory,
                best_hp_params=best_hp_params,
                best_hardcoded_params=estimator_hardcoded_params,
                default_params=default_params,
            )
        elif len(estimator_hardcoded_params) > 0:
            estimator = estimator_fn(estimator_hardcoded_params)
            all_estimator_params = estimator.get_params()
            default_params_keys = all_estimator_params.keys() - estimator_hardcoded_params.keys()
            default_params = {k: all_estimator_params[k] for k in default_params_keys}
            self._write_best_parameters_outputs(
                output_directory,
                best_hardcoded_params=estimator_hardcoded_params,
                default_params=default_params,
            )
        else:
            estimator = estimator_fn()
            default_params = estimator.get_params()
            self._write_best_parameters_outputs(output_directory, default_params=default_params)
        return estimator

    def _resolve_estimator_plugin(self, plugin_str, X_train, y_train, output_directory):
        plugin_str = plugin_str.replace("/", ".")
        estimator_fn = importlib.import_module(
            f"mlflow.recipes.steps.{plugin_str}"
        ).get_estimator_and_best_params
        estimator, best_parameters = estimator_fn(
            X_train,
            y_train,
            self.task,
            self.extended_task,
            self.step_config,
            self.recipe_root,
            self.evaluation_metrics,
            self.primary_metric,
        )
        self.best_estimator_name = estimator.__class__.__name__
        self.best_estimator_class = (
            f"{estimator.__class__.__module__}.{estimator.__class__.__name__}"
        )
        self.best_parameters = best_parameters
        self._write_best_parameters_outputs(output_directory, automl_params=best_parameters)
        return estimator

    def _resolve_estimator(self, X_train, y_train, validation_df, run, output_directory):
        using_plugin = self.step_config.get("using", "custom")
        if using_plugin == "custom":
            return self._get_user_defined_estimator(
                X_train, y_train, validation_df, run, output_directory
            )
        else:
            return self._resolve_estimator_plugin(using_plugin, X_train, y_train, output_directory)

    def _get_leaderboard_df(self, run, eval_metrics):
        import pandas as pd

        mlflow_client = MlflowClient()
        exp_id = _get_experiment_id()

        primary_metric_greater_is_better = self.evaluation_metrics[
            self.primary_metric
        ].greater_is_better
        primary_metric_order = "DESC" if primary_metric_greater_is_better else "ASC"

        search_max_results = 100
        search_result = mlflow_client.search_runs(
            experiment_ids=exp_id,
            run_view_type=ViewType.ACTIVE_ONLY,
            max_results=search_max_results,
            order_by=[f"metrics.{self.primary_metric} {primary_metric_order}"],
        )

        metric_names = self.evaluation_metrics.keys()

        leaderboard_items = []
        for old_run in search_result:
            if all(metric_key in old_run.data.metrics for metric_key in metric_names):
                leaderboard_items.append(
                    {
                        "Run ID": old_run.info.run_id,
                        "Run Time": datetime.datetime.fromtimestamp(
                            old_run.info.start_time // 1000
                        ),
                        **{
                            metric_name: old_run.data.metrics[metric_key]
                            for metric_name, metric_key in zip(metric_names, metric_names)
                        },
                    }
                )

        top_leaderboard_items = [
            {"Model Rank": i + 1, **t} for i, t in enumerate(leaderboard_items[:2])
        ]

        if (
            len(top_leaderboard_items) == 2
            and top_leaderboard_items[0][self.primary_metric]
            == top_leaderboard_items[1][self.primary_metric]
        ):
            # If top1 and top2 model primary metrics are equal,
            # then their rank are both 1.
            top_leaderboard_items[1]["Model Rank"] = "1"

        top_leaderboard_item_index_values = ["Best", "2nd Best"][: len(top_leaderboard_items)]

        latest_model_item = {
            "Run ID": run.info.run_id,
            "Run Time": datetime.datetime.fromtimestamp(run.info.start_time // 1000),
            **eval_metrics["validation"],
        }

        for i, leaderboard_item in enumerate(leaderboard_items):
            latest_value = latest_model_item[self.primary_metric]
            historical_value = leaderboard_item[self.primary_metric]
            if (primary_metric_greater_is_better and latest_value >= historical_value) or (
                not primary_metric_greater_is_better and latest_value <= historical_value
            ):
                latest_model_item["Model Rank"] = str(i + 1)
                break
        else:
            latest_model_item["Model Rank"] = f"> {len(leaderboard_items)}"

        # metric columns order: primary metric, custom metrics, builtin metrics.
        def sorter(m):
            if m == self.primary_metric:
                return 0, m
            elif self.evaluation_metrics[m].custom_function is not None:
                return 1, m
            else:
                return 2, m

        metric_columns = sorted(metric_names, key=sorter)

        return (
            pd.DataFrame.from_records(
                [latest_model_item, *top_leaderboard_items],
                columns=["Model Rank", *metric_columns, "Run Time", "Run ID"],
            )
            .apply(
                lambda s: s.map(lambda x: f"{x:.6g}") if s.name in metric_names else s,
                axis=0,
            )
            .set_axis(["Latest"] + top_leaderboard_item_index_values, axis="index")
            .transpose()
        )

    def _get_tuning_df(self, run, params=None):
        exp_id = _get_experiment_id()
        primary_metric_tag = f"metrics.{self.primary_metric}"
        order_str = (
            "DESC" if self.evaluation_metrics_greater_is_better[self.primary_metric] else "ASC"
        )
        tuning_runs = mlflow.search_runs(
            [exp_id],
            filter_string=f"tags.mlflow.parentRunId like '{run.info.run_id}'",
            order_by=[f"{primary_metric_tag} {order_str}", "attribute.start_time ASC"],
        )
        if params:
            params = [f"params.{param}" for param in params]
            tuning_runs = tuning_runs.filter([f"metrics.{self.primary_metric}", *params])
        else:
            tuning_runs = tuning_runs.filter([f"metrics.{self.primary_metric}"])
        tuning_runs = tuning_runs.reset_index().rename(
            columns={"index": "Model Rank", primary_metric_tag: self.primary_metric}
        )
        tuning_runs["Model Rank"] += 1
        return tuning_runs.head(10)

    def _build_step_card(
        self,
        eval_metrics,
        pred_and_error_df,
        model_schema,
        run_id,
        model_uri,
        worst_examples_df,
        train_df,
        output_directory,
        leaderboard_df=None,
        tuning_df=None,
        calibrated_plot=None,
    ):
        import pandas as pd
        from sklearn import set_config
        from sklearn.utils import estimator_html_repr

        card = BaseCard(self.recipe_name, self.name)
        # Tab 0: model performance summary.
        metric_df = (
            get_merged_eval_metrics(
                eval_metrics,
                ordered_metric_names=[
                    self.primary_metric,
                    *self.user_defined_custom_metrics.keys(),
                ],
            )
            .reset_index()
            .rename(columns={"index": "Metric"})
        )

        def row_style(row):
            if row.Metric == self.primary_metric or row.Metric in self.user_defined_custom_metrics:
                return pd.Series("font-weight: bold", row.index)
            else:
                return pd.Series("", row.index)

        metric_table_html = BaseCard.render_table(
            metric_df.style.format({"training": "{:.6g}", "validation": "{:.6g}"}).apply(
                row_style, axis=1
            )
        )

        # Tab 1: Model performance.
        card.add_tab(
            "Model Performance",
            "<h3 class='section-title'>Summary Metrics (Validation)</h3>{{ METRICS }} ",
        ).add_html("METRICS", metric_table_html)

        if not self.skip_data_profiling:
            # Tab 2: Prediction and error data profile.
            pred_and_error_df_profile = get_pandas_data_profiles(
                [
                    [
                        "Predictions and Errors (Validation Dataset)",
                        pred_and_error_df.reset_index(drop=True),
                    ]
                ]
            )
            card.add_tab("Data Profile (Predictions)", "{{PROFILE}}").add_pandas_profile(
                "PROFILE", pred_and_error_df_profile
            )
        # Tab 3: Model architecture.
        set_config(display="diagram")
        model_repr = estimator_html_repr(
            mlflow.sklearn.load_model(
                os.path.join(model_uri, "artifacts", TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH)
            )
        )
        card.add_tab("Model Architecture", "{{MODEL_ARCH}}").add_html("MODEL_ARCH", model_repr)

        # Tab 4: Inferred model (transformer + estimator) schema.
        def render_schema(inputs, title):
            from mlflow.types import ColSpec

            table = BaseCard.render_table(
                {
                    "Name": "  " + (spec.name or "-"),
                    "Type": repr(spec.type) if isinstance(spec, ColSpec) else repr(spec),
                }
                for spec in inputs
            )
            return f'<div style="margin: 5px"><h2>{title}</h2>{table}</div>'

        schema_tables = [render_schema(model_schema.inputs.inputs, "Inputs")]
        if model_schema.outputs:
            schema_tables += [render_schema(model_schema.outputs.inputs, "Outputs")]

        card.add_tab("Model Schema", "{{MODEL_SCHEMA}}").add_html(
            "MODEL_SCHEMA",
            '<div style="display: flex">{tables}</div>'.format(tables="\n".join(schema_tables)),
        )

        # Tab 5: Examples with Largest Prediction Error
        if not worst_examples_df.empty:
            (
                card.add_tab("Worst Predictions", "{{ WORST_EXAMPLES_TABLE }}").add_html(
                    "WORST_EXAMPLES_TABLE", BaseCard.render_table(worst_examples_df)
                )
            )

        if calibrated_plot:
            calibrated_plot_location = os.path.join(output_directory, "calibrated_plot_location")
            calibrated_plot.figure_.savefig(calibrated_plot_location, format="png")
            (
                card.add_tab("Prob. Calibration", "{{ CALIBRATED_PLOT }}").add_image(
                    "CALIBRATED_PLOT", calibrated_plot_location
                )
            )

        # Tab 6: Worst predictions profile vs train profile.
        if not self.skip_data_profiling:
            worst_prediction_profile = get_pandas_data_profiles(
                [
                    ["Worst Predictions", worst_examples_df.reset_index(drop=True)],
                    ["Train", train_df.reset_index(drop=True)],
                ]
            )
            card.add_tab(
                "Data Profile (Worst vs Train)", "{{ WORST_EXAMPLES_COMP }}"
            ).add_pandas_profile("WORST_EXAMPLES_COMP", worst_prediction_profile)

        # Tab 7: Leaderboard
        if leaderboard_df is not None:
            (
                card.add_tab("Leaderboard", "{{ LEADERBOARD_TABLE }}").add_html(
                    "LEADERBOARD_TABLE", BaseCard.render_table(leaderboard_df, hide_index=False)
                )
            )

        # Tab 8: Best Parameters (AutoML and Tuning)
        is_automl_run = self.step_config["using"].startswith("automl")
        best_parameters_yaml = os.path.join(output_directory, "best_parameters.yaml")

        if os.path.exists(best_parameters_yaml):
            best_parameters_card_tab = card.add_tab(
                f"Best Parameters {' (AutoML)' if is_automl_run else ''}",
                "{{ BEST_PARAMETERS }} ",
            )

            if is_automl_run:
                automl_estimator_str = (
                    f"<b>Best estimator:</b><br>"
                    f"<pre>{self.best_estimator_name}</pre><br>"
                    f"<b>Best estimator class:</b><br>"
                    f"<pre>{self.best_estimator_class}</pre><br><br>"
                )
            else:
                automl_estimator_str = ""

            with open(best_parameters_yaml) as f:
                best_parameters = f.read()
            best_parameters_card_tab.add_html(
                "BEST_PARAMETERS",
                f"{automl_estimator_str}<b>Best parameters:</b><br>"
                f"<pre>{best_parameters}</pre><br><br>",
            )

        # Tab 9: HP trials
        if tuning_df is not None:
            tuning_trials_card_tab = card.add_tab(
                "Tuning Trials",
                "{{ SEARCH_SPACE }}" + "{{ TUNING_TABLE_TITLE }}" + "{{ TUNING_TABLE }}",
            )
            tuning_params = yaml.dump(self.step_config["tuning"]["parameters"])
            tuning_trials_card_tab.add_html(
                "SEARCH_SPACE",
                f"<b>Tuning search space:</b> <br><pre>{tuning_params}</pre><br><br>",
            )
            tuning_trials_card_tab.add_html("TUNING_TABLE_TITLE", "<b>Tuning results:</b><br>")
            tuning_trials_card_tab.add_html(
                "TUNING_TABLE",
                BaseCard.render_table(
                    tuning_df.style.apply(
                        lambda row: pd.Series("font-weight: bold", row.index),
                        axis=1,
                        subset=tuning_df.index[0],
                    ),
                    hide_index=True,
                ),
            )

        # Tab 10: Warning log outputs.
        warning_output_path = os.path.join(output_directory, "warning_logs.txt")
        if os.path.exists(warning_output_path):
            warnings_output_tab = card.add_tab("Warning Logs", "{{ STEP_WARNINGS }}")
            with open(warning_output_path) as f:
                warnings_output_tab.add_html("STEP_WARNINGS", f"<pre>{f.read()}</pre>")

        # Tab 11: Run summary.
        run_card_tab = card.add_tab(
            "Run Summary",
            "{{ RUN_ID }} " + "{{ MODEL_URI }}" + "{{ EXE_DURATION }}" + "{{ LAST_UPDATE_TIME }}",
        )
        model_uri_path = f"runs:/{run_id}/train/model"
        run_url = get_databricks_run_url(
            tracking_uri=mlflow.get_tracking_uri(),
            run_id=run_id,
        )
        model_url = get_databricks_run_url(
            tracking_uri=mlflow.get_tracking_uri(),
            run_id=run_id,
            artifact_path=re.sub(rf"^.*?{run_id}", "", model_uri_path),
        )

        if run_url is not None:
            run_card_tab.add_html(
                "RUN_ID", f"<b>MLflow Run ID:</b> <a href={run_url}>{run_id}</a><br><br>"
            )
        else:
            run_card_tab.add_markdown("RUN_ID", f"**MLflow Run ID:** `{run_id}`")

        if model_url is not None:
            run_card_tab.add_html(
                "MODEL_URI", f"<b>MLflow Model URI:</b> <a href={model_url}>{model_uri_path}</a>"
            )
        else:
            run_card_tab.add_markdown("MODEL_URI", f"**MLflow Model URI:** `{model_uri_path}`")

        return card

    @classmethod
    def from_recipe_config(cls, recipe_config, recipe_root):
        step_config = {}
        if recipe_config.get("steps", {}).get("train", {}) is not None:
            step_config.update(recipe_config.get("steps", {}).get("train", {}))
        if recipe_config.get("custom_metrics") is not None:
            step_config["custom_metrics"] = recipe_config["custom_metrics"]
        if recipe_config.get("primary_metric") is not None:
            step_config["primary_metric"] = recipe_config["primary_metric"]
        step_config["recipe"] = recipe_config.get("recipe")
        step_config["profile"] = recipe_config.get("profile")
        step_config["target_col"] = recipe_config.get("target_col")
        if "positive_class" in recipe_config:
            step_config["positive_class"] = recipe_config.get("positive_class")
        step_config.update(
            get_recipe_tracking_config(
                recipe_root_path=recipe_root,
                recipe_config=recipe_config,
            ).to_dict()
        )
        return cls(step_config, recipe_root, recipe_config=recipe_config)

    @property
    def name(self):
        return "train"

    @property
    def environment(self):
        environ = get_databricks_env_vars(tracking_uri=self.tracking_config.tracking_uri)
        environ.update(get_run_tags_env_vars(recipe_root_path=self.recipe_root))
        return environ

    def get_artifacts(self):
        return [
            ModelArtifact("model", self.recipe_root, self.name, self.tracking_config.tracking_uri),
            RunArtifact("run", self.recipe_root, self.name, self.tracking_config.tracking_uri),
            HyperParametersArtifact("best_parameters", self.recipe_root, self.name),
            DataframeArtifact(
                "predicted_training_data",
                self.recipe_root,
                self.name,
                TrainStep.PREDICTED_TRAINING_DATA_RELATIVE_PATH,
            ),
        ]

    def step_class(self):
        return StepClass.TRAINING

    def _tune_and_get_best_estimator_params(
        self,
        parent_run_id,
        estimator_hardcoded_params,
        estimator_fn,
        X_train,
        y_train,
        validation_df,
    ):
        tuning_params = self.step_config["tuning"]
        try:
            from hyperopt import SparkTrials, Trials, fmin, space_eval
        except ModuleNotFoundError:
            raise MlflowException(
                "Hyperopt not installed and is required if tuning is enabled",
                error_code=BAD_REQUEST,
            )

        # wrap training in objective fn
        def objective(X_train, y_train, validation_df, hyperparameter_args, on_worker=False):
            run_name = self.tracking_config.run_name
            if on_worker:
                client = MlflowClient()
                parent_tags = client.get_run(parent_run_id).data.tags
                child_run = client.create_run(
                    _get_experiment_id(),
                    tags={**parent_tags, "mlflow.parentRunId": parent_run_id},
                    run_name=run_name,
                )
                run_args = {"run_id": child_run.info.run_id}
            else:
                run_args = {"run_name": run_name, "nested": True}
            with mlflow.start_run(**run_args) as tuning_run:
                estimator_args = dict(estimator_hardcoded_params, **hyperparameter_args)
                estimator = estimator_fn(estimator_args)

                sample_fraction = self.step_config["sample_fraction"]

                # if sparktrials, then read from broadcast
                if tuning_params["parallelism"] > 1:
                    X_train = X_train.value
                    y_train = y_train.value
                    validation_df = validation_df.value

                X_train_sampled = X_train.sample(frac=sample_fraction, random_state=42)
                y_train_sampled = y_train.sample(frac=sample_fraction, random_state=42)

                fitted_estimator, _ = self._fitted_estimator(
                    estimator, X_train_sampled, y_train_sampled
                )

                logged_estimator = self._log_estimator_to_mlflow(
                    fitted_estimator, X_train_sampled, on_worker=on_worker
                )

                eval_result = mlflow.evaluate(
                    model=logged_estimator.model_uri,
                    data=validation_df,
                    targets=self.target_col,
                    model_type=_get_model_type_from_template(self.recipe),
                    evaluators="default",
                    extra_metrics=_load_custom_metrics(
                        self.recipe_root,
                        self.evaluation_metrics.values(),
                    ),
                    evaluator_config={
                        "log_model_explainability": False,
                        "pos_label": self.positive_class,
                    },
                )
                autologged_params = mlflow.get_run(run_id=tuning_run.info.run_id).data.params
                manual_log_params = {}
                for param_name, param_value in estimator_args.items():
                    if param_name in autologged_params:
                        if not TrainStep.is_tuning_param_equal(
                            param_value, autologged_params[param_name]
                        ):
                            _logger.warning(
                                f"Failed to log search space parameter due to conflict. "
                                f"Attempted to log search space parameter {param_name} as "
                                f"{param_value}, but {param_name} is already logged as "
                                f"{autologged_params[param_name]} during training. "
                                f"Are you passing `estimator_params` properly to the "
                                f" estimator?"
                            )
                    else:
                        manual_log_params[param_name] = param_value

                if len(manual_log_params) > 0:
                    mlflow.log_params(manual_log_params)

                # return +/- metric
                transformed_metrics = transform_multiclass_metrics_dict(
                    eval_result.metrics, self.extended_task
                )
                sign = -1 if self.evaluation_metrics_greater_is_better[self.primary_metric] else 1
                return sign * transformed_metrics[self.primary_metric]

        search_space = TrainStep.construct_search_space_from_yaml(tuning_params["parameters"])
        algo_type, algo_name = tuning_params["algorithm"].rsplit(".", 1)
        tuning_algo = getattr(importlib.import_module(algo_type, "hyperopt"), algo_name)
        max_trials = tuning_params["max_trials"]
        parallelism = tuning_params["parallelism"]

        if parallelism > 1:
            from pyspark.sql import SparkSession

            spark_session = SparkSession.builder.config(
                "spark.databricks.mlflow.trackHyperopt.enabled", "false"
            ).getOrCreate()
            sc = spark_session.sparkContext

            X_train = sc.broadcast(X_train)
            y_train = sc.broadcast(y_train)
            validation_df = sc.broadcast(validation_df)

            hp_trials = SparkTrials(parallelism, spark_session=spark_session)
            on_worker = True
        else:
            hp_trials = Trials()
            on_worker = False

        fmin_kwargs = {
            "fn": lambda params: objective(
                X_train, y_train, validation_df, params, on_worker=on_worker
            ),
            "space": search_space,
            "algo": tuning_algo,
            "max_evals": max_trials,
            "trials": hp_trials,
        }
        if "early_stop_fn" in tuning_params:
            try:
                train_module_name, early_stop_fn_name = tuning_params["early_stop_fn"].rsplit(
                    ".", 1
                )
            except ValueError:
                early_stop_fn_name = tuning_params["early_stop_fn"]
                train_module_name = _USER_DEFINED_TRAIN_STEP_MODULE
            early_stop_fn = getattr(importlib.import_module(train_module_name), early_stop_fn_name)
            fmin_kwargs["early_stop_fn"] = early_stop_fn
        best_hp_params = fmin(**fmin_kwargs)
        best_hp_params = space_eval(search_space, best_hp_params)
        best_hp_estimator_loss = hp_trials.best_trial["result"]["loss"]
        if len(estimator_hardcoded_params) > 1:
            hardcoded_estimator_loss = objective(
                X_train, y_train, validation_df, estimator_hardcoded_params
            )

            if best_hp_estimator_loss < hardcoded_estimator_loss:
                best_hardcoded_params = {
                    param_name: param_value
                    for param_name, param_value in estimator_hardcoded_params.items()
                    if param_name not in best_hp_params
                }
            else:
                best_hp_params = {}
                best_hardcoded_params = estimator_hardcoded_params
        else:
            best_hardcoded_params = {}
        return (best_hardcoded_params, best_hp_params)

    def _log_estimator_to_mlflow(self, estimator, X_train_sampled, on_worker=False):
        from mlflow.models import infer_signature

        if hasattr(estimator, "best_score_") and (type(estimator.best_score_) in [int, float]):
            mlflow.log_metric("best_cv_score", estimator.best_score_)
        if hasattr(estimator, "best_params_"):
            mlflow.log_params(estimator.best_params_)

        if on_worker:
            mlflow.log_params(estimator.get_params())
            estimator_tags = {
                "estimator_name": estimator.__class__.__name__,
                "estimator_class": (
                    estimator.__class__.__module__ + "." + estimator.__class__.__name__
                ),
            }
            mlflow.set_tags(estimator_tags)
        estimator_schema = infer_signature(
            X_train_sampled, estimator.predict(X_train_sampled.copy())
        )
        return mlflow.sklearn.log_model(
            estimator,
            f"{self.name}/estimator",
            signature=estimator_schema,
            code_paths=self.code_paths,
        )

    def _write_best_parameters_outputs(
        self,
        output_directory,
        best_hp_params=None,
        best_hardcoded_params=None,
        automl_params=None,
        default_params=None,
    ):
        if best_hp_params or best_hardcoded_params or automl_params or default_params:
            best_parameters_path = os.path.join(output_directory, "best_parameters.yaml")
            if os.path.exists(best_parameters_path):
                os.remove(best_parameters_path)
            with open(best_parameters_path, "a") as file:
                self._write_one_param_output(automl_params or {}, file, "automl parameters")
                self._write_one_param_output(best_hp_params or {}, file, "tuned hyperparameters")
                self._write_one_param_output(
                    best_hardcoded_params or {}, file, "hardcoded parameters"
                )
                self._write_one_param_output(default_params or {}, file, "default parameters")
            mlflow.log_artifact(best_parameters_path, artifact_path="train")

    def _write_one_param_output(self, params, file, caption):
        if params:
            file.write(f"# {caption} \n")
            self._safe_dump_with_numeric_values(params, file, default_flow_style=False)
            file.write("\n")

    def _safe_dump_with_numeric_values(self, data, file, **kwargs):
        import numpy as np

        processed_data = {}
        for key, value in data.items():
            if isinstance(value, np.floating):
                processed_data[key] = float(value)
            elif isinstance(value, np.integer):
                processed_data[key] = int(value)
            elif isinstance(value, dict):
                processed_data[key] = str(value)
            else:
                processed_data[key] = str(value)

        if len(processed_data) > 0:
            yaml.safe_dump(processed_data, file, **kwargs)

    def _rebalance_classes(self, train_df):
        import pandas as pd

        resampling_minority_percentage = self.step_config.get(
            "resampling_minority_percentage", _REBALANCING_DEFAULT_RATIO
        )

        df_positive_class = train_df[train_df[self.target_col] == self.positive_class]
        df_negative_class = train_df[train_df[self.target_col] != self.positive_class]

        if len(df_positive_class) > len(df_negative_class):
            df_minority_class, df_majority_class = df_negative_class, df_positive_class
        else:
            df_minority_class, df_majority_class = df_positive_class, df_negative_class

        original_minority_percentage = len(df_minority_class) / len(train_df)
        if original_minority_percentage >= resampling_minority_percentage:
            _logger.info(
                f"Class imbalance of {original_minority_percentage:.2f} "
                f"is better than {resampling_minority_percentage}, no need to rebalance"
            )
            return train_df

        _logger.info(
            f"Detected class imbalance: minority class percentage is "
            f"{original_minority_percentage:.2f}"
        )

        majority_class_target = int(
            len(df_minority_class)
            * (1 - resampling_minority_percentage)
            / resampling_minority_percentage
        )
        df_majority_downsampled = df_majority_class.sample(majority_class_target)
        _logger.info(
            f"After downsampling: minority class percentage is {resampling_minority_percentage:.2f}"
        )
        return pd.concat([df_minority_class, df_majority_downsampled], axis=0).sample(frac=1)
