import functools
import mlflow
from mlflow import MlflowClient
from mlflow.exceptions import MlflowException
from mlflow.models.evaluation.base import (
    ModelEvaluator,
    EvaluationResult,
)
from mlflow.entities.metric import Metric
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
from mlflow.utils.file_utils import TempDir
from mlflow.models.utils import plot_lines
from mlflow.models.evaluation.artifacts import (
    ImageEvaluationArtifact,
    CsvEvaluationArtifact,
    NumpyEvaluationArtifact,
    _infer_artifact_type_and_ext,
    JsonEvaluationArtifact,
)
from mlflow.pyfunc import _ServedPyFuncModel
from mlflow.utils.proto_json_utils import NumpyEncoder
from mlflow.utils.time_utils import get_current_time_millis

from sklearn import metrics as sk_metrics
from sklearn.pipeline import Pipeline as sk_Pipeline
import math
import json
from collections import namedtuple
from typing import NamedTuple, Callable
import tempfile
import pandas as pd
import numpy as np
import copy
import shutil
import pickle
from functools import partial
import logging
from packaging.version import Version
import pathlib

_logger = logging.getLogger(__name__)

_DEFAULT_SAMPLE_ROWS_FOR_SHAP = 2000


def _is_categorical(values):
    """
    Infer whether input values are categorical on best effort.
    Return True represent they are categorical, return False represent we cannot determine result.
    """
    dtype_name = pd.Series(values).convert_dtypes().dtype.name.lower()
    return dtype_name in ["category", "string", "boolean"]


def _is_continuous(values):
    """
    Infer whether input values is continuous on best effort.
    Return True represent they are continous, return False represent we cannot determine result.
    """
    dtype_name = pd.Series(values).convert_dtypes().dtype.name.lower()
    return dtype_name.startswith("float")


def _infer_model_type_by_labels(labels):
    """
    Infer model type by target values.
    """
    if _is_categorical(labels):
        return "classifier"
    elif _is_continuous(labels):
        return "regressor"
    else:
        return None  # Unknown


def _extract_raw_model(model):
    """
    Return a tuple of (model_loader_module, raw_model)
    """
    model_loader_module = model.metadata.flavors["python_function"]["loader_module"]
    try:
        if model_loader_module == "mlflow.sklearn" and not isinstance(model, _ServedPyFuncModel):
            raw_model = model._model_impl
        else:
            raw_model = None
    except Exception as e:
        _logger.warning(
            f"Raw model resolution fails unexpectedly on PyFuncModel {model!r}, "
            f"error message is {e}"
        )
        raw_model = None

    return model_loader_module, raw_model


def _extract_predict_fn(model, raw_model):
    predict_fn = model.predict
    predict_proba_fn = None

    if raw_model is not None:
        predict_fn = raw_model.predict
        predict_proba_fn = getattr(raw_model, "predict_proba", None)

        try:
            import xgboost

            if isinstance(raw_model, xgboost.XGBModel):
                # Because shap evaluation will pass evaluation data in ndarray format
                # (without feature names), if set validate_features=True it will raise error.
                predict_fn = partial(predict_fn, validate_features=False)
                if predict_proba_fn is not None:
                    predict_proba_fn = partial(predict_proba_fn, validate_features=False)
        except ImportError:
            pass

    return predict_fn, predict_proba_fn


def _get_regressor_metrics(y, y_pred, sample_weights):
    sum_on_target = (
        (np.array(y) * np.array(sample_weights)).sum() if sample_weights is not None else sum(y)
    )
    return {
        "example_count": len(y),
        "mean_absolute_error": sk_metrics.mean_absolute_error(
            y, y_pred, sample_weight=sample_weights
        ),
        "mean_squared_error": sk_metrics.mean_squared_error(
            y, y_pred, sample_weight=sample_weights
        ),
        "root_mean_squared_error": sk_metrics.mean_squared_error(
            y, y_pred, sample_weight=sample_weights, squared=False
        ),
        "sum_on_target": sum_on_target,
        "mean_on_target": sum_on_target / len(y),
        "r2_score": sk_metrics.r2_score(y, y_pred, sample_weight=sample_weights),
        "max_error": sk_metrics.max_error(y, y_pred),
        "mean_absolute_percentage_error": sk_metrics.mean_absolute_percentage_error(
            y, y_pred, sample_weight=sample_weights
        ),
    }


def _get_binary_sum_up_label_pred_prob(positive_class_index, positive_class, y, y_pred, y_probs):
    y = np.array(y)
    y_bin = np.where(y == positive_class, 1, 0)
    y_pred_bin = None
    y_prob_bin = None
    if y_pred is not None:
        y_pred = np.array(y_pred)
        y_pred_bin = np.where(y_pred == positive_class, 1, 0)

    if y_probs is not None:
        y_probs = np.array(y_probs)
        y_prob_bin = y_probs[:, positive_class_index]

    return y_bin, y_pred_bin, y_prob_bin


def _get_common_classifier_metrics(
    *, y_true, y_pred, y_proba, labels, average, pos_label, sample_weights
):
    metrics = {
        "example_count": len(y_true),
        "accuracy_score": sk_metrics.accuracy_score(y_true, y_pred, sample_weight=sample_weights),
        "recall_score": sk_metrics.recall_score(
            y_true,
            y_pred,
            average=average,
            pos_label=pos_label,
            sample_weight=sample_weights,
        ),
        "precision_score": sk_metrics.precision_score(
            y_true,
            y_pred,
            average=average,
            pos_label=pos_label,
            sample_weight=sample_weights,
        ),
        "f1_score": sk_metrics.f1_score(
            y_true,
            y_pred,
            average=average,
            pos_label=pos_label,
            sample_weight=sample_weights,
        ),
    }
    if y_proba is not None:
        metrics["log_loss"] = sk_metrics.log_loss(
            y_true, y_proba, labels=labels, sample_weight=sample_weights
        )

    return metrics


def _get_binary_classifier_metrics(
    *, y_true, y_pred, y_proba=None, labels=None, pos_label=1, sample_weights=None
):
    tn, fp, fn, tp = sk_metrics.confusion_matrix(y_true, y_pred).ravel()
    return {
        "true_negatives": tn,
        "false_positives": fp,
        "false_negatives": fn,
        "true_positives": tp,
        **_get_common_classifier_metrics(
            y_true=y_true,
            y_pred=y_pred,
            y_proba=y_proba,
            labels=labels,
            average="binary",
            pos_label=pos_label,
            sample_weights=sample_weights,
        ),
    }


def _get_multiclass_classifier_metrics(
    *,
    y_true,
    y_pred,
    y_proba=None,
    labels=None,
    average="weighted",
    sample_weights=None,
):
    metrics = _get_common_classifier_metrics(
        y_true=y_true,
        y_pred=y_pred,
        y_proba=y_proba,
        labels=labels,
        average=average,
        pos_label=None,
        sample_weights=sample_weights,
    )
    if average in ("macro", "weighted") and y_proba is not None:
        metrics.update(
            roc_auc=sk_metrics.roc_auc_score(
                y_true=y_true,
                y_score=y_proba,
                sample_weight=sample_weights,
                average=average,
                multi_class="ovr",
            )
        )
    return metrics


def _get_classifier_per_class_metrics_collection_df(y, y_pred, labels, sample_weights):
    per_class_metrics_list = []
    for positive_class_index, positive_class in enumerate(labels):
        (
            y_bin,
            y_pred_bin,
            _,
        ) = _get_binary_sum_up_label_pred_prob(
            positive_class_index, positive_class, y, y_pred, None
        )
        per_class_metrics = {"positive_class": positive_class}
        per_class_metrics.update(
            _get_binary_classifier_metrics(
                y_true=y_bin,
                y_pred=y_pred_bin,
                pos_label=1,
                sample_weights=sample_weights,
            )
        )
        per_class_metrics_list.append(per_class_metrics)

    return pd.DataFrame(per_class_metrics_list)


def _get_dataframe_with_renamed_columns(x, new_column_names):
    """
    Downstream inference functions may expect a pd.DataFrame to be created from x. However,
    if x is already a pd.DataFrame, and new_column_names != x.columns, we cannot simply call
    pd.DataFrame(x, columns=new_column_names) because the resulting pd.DataFrame will contain
    NaNs for every column in new_column_names that does not exist in x.columns. This function
    instead creates a new pd.DataFrame object from x, and then explicitly renames the columns
    to avoid NaNs.

    :param x: :param data: A data object, such as a Pandas DataFrame, numPy array, or list
    :param new_column_names: Column names for the output Pandas DataFrame
    :return: A pd.DataFrame with x as data, with columns new_column_names
    """
    df = pd.DataFrame(x)
    return df.rename(columns=dict(zip(df.columns, new_column_names)))


_Curve = namedtuple("_Curve", ["plot_fn", "plot_fn_args", "auc"])


def _gen_classifier_curve(
    is_binomial,
    y,
    y_probs,
    labels,
    pos_label,
    curve_type,
    sample_weights,
):
    """
    Generate precision-recall curve or ROC curve for classifier.
    :param is_binomial: True if it is binary classifier otherwise False
    :param y: True label values
    :param y_probs: if binary classifier, the predicted probability for positive class.
                    if multiclass classifier, the predicted probabilities for all classes.
    :param labels: The set of labels.
    :param pos_label: The label of the positive class.
    :param curve_type: "pr" or "roc"
    :param sample_weights: Optional sample weights.
    :return: An instance of "_Curve" which includes attributes "plot_fn", "plot_fn_args", "auc".
    """
    if curve_type == "roc":

        def gen_line_x_y_label_auc(_y, _y_prob, _pos_label):
            fpr, tpr, _ = sk_metrics.roc_curve(
                _y,
                _y_prob,
                sample_weight=sample_weights,
                # For multiclass classification where a one-vs-rest ROC curve is produced for each
                # class, the positive label is binarized and should not be included in the plot
                # legend
                pos_label=_pos_label if _pos_label == pos_label else None,
            )
            auc = sk_metrics.roc_auc_score(y_true=_y, y_score=_y_prob, sample_weight=sample_weights)
            return fpr, tpr, f"AUC={auc:.3f}", auc

        xlabel = "False Positive Rate"
        ylabel = "True Positive Rate"
        title = "ROC curve"
        if pos_label:
            xlabel = f"False Positive Rate (Positive label: {pos_label})"
            ylabel = f"True Positive Rate (Positive label: {pos_label})"
    elif curve_type == "pr":

        def gen_line_x_y_label_auc(_y, _y_prob, _pos_label):
            precision, recall, _ = sk_metrics.precision_recall_curve(
                _y,
                _y_prob,
                sample_weight=sample_weights,
                # For multiclass classification where a one-vs-rest precision-recall curve is
                # produced for each class, the positive label is binarized and should not be
                # included in the plot legend
                pos_label=_pos_label if _pos_label == pos_label else None,
            )
            # NB: We return average precision score (AP) instead of AUC because AP is more
            # appropriate for summarizing a precision-recall curve
            ap = sk_metrics.average_precision_score(
                y_true=_y, y_score=_y_prob, pos_label=_pos_label, sample_weight=sample_weights
            )
            return recall, precision, f"AP={ap:.3f}", ap

        xlabel = "Recall"
        ylabel = "Precision"
        title = "Precision recall curve"
        if pos_label:
            xlabel = f"Recall (Positive label: {pos_label})"
            ylabel = f"Precision (Positive label: {pos_label})"
    else:
        assert False, "illegal curve type"

    if is_binomial:
        x_data, y_data, line_label, auc = gen_line_x_y_label_auc(y, y_probs, pos_label)
        data_series = [(line_label, x_data, y_data)]
    else:
        curve_list = []
        for positive_class_index, positive_class in enumerate(labels):
            y_bin, _, y_prob_bin = _get_binary_sum_up_label_pred_prob(
                positive_class_index, positive_class, y, labels, y_probs
            )

            x_data, y_data, line_label, auc = gen_line_x_y_label_auc(
                y_bin, y_prob_bin, _pos_label=1
            )
            curve_list.append((positive_class, x_data, y_data, line_label, auc))

        data_series = [
            (f"label={positive_class},{line_label}", x_data, y_data)
            for positive_class, x_data, y_data, line_label, _ in curve_list
        ]
        auc = [auc for _, _, _, _, auc in curve_list]

    def _do_plot(**kwargs):
        from matplotlib import pyplot

        _, ax = plot_lines(**kwargs)
        dash_line_args = {
            "color": "gray",
            "alpha": 0.3,
            "drawstyle": "default",
            "linestyle": "dashed",
        }
        if curve_type == "pr":
            ax.plot([0, 1], [1, 0], **dash_line_args)
        elif curve_type == "roc":
            ax.plot([0, 1], [0, 1], **dash_line_args)

        if is_binomial:
            ax.legend(loc="best")
        else:
            ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
            pyplot.subplots_adjust(right=0.6, bottom=0.25)

    return _Curve(
        plot_fn=_do_plot,
        plot_fn_args={
            "data_series": data_series,
            "xlabel": xlabel,
            "ylabel": ylabel,
            "line_kwargs": {"drawstyle": "steps-post", "linewidth": 1},
            "title": title,
        },
        auc=auc,
    )


_matplotlib_config = {
    "figure.dpi": 175,
    "figure.figsize": [6.0, 4.0],
    "figure.autolayout": True,
    "font.size": 8,
}


class _CustomMetric(NamedTuple):
    """
    A namedtuple representing a custom metric function and its properties.

    function : the custom metric function
    name : the name of the custom metric function
    index : the index of the function in the ``custom_metrics`` argument of mlflow.evaluate
    """

    function: Callable
    name: str
    index: int


class _CustomArtifact(NamedTuple):
    """
    A namedtuple representing a custom artifact function and its properties.

    function : the custom artifact function
    name : the name of the custom artifact function
    index : the index of the function in the ``custom_artifacts`` argument of mlflow.evaluate
    artifacts_dir : the path to a temporary directory to store produced artifacts of the function
    """

    function: Callable
    name: str
    index: int
    artifacts_dir: str


def _is_numeric(value):
    return isinstance(value, (int, float, np.number))


def _evaluate_custom_metric(custom_metric_tuple, eval_df, builtin_metrics):
    """
    This function calls the `custom_metric` function and performs validations on the returned
    result to ensure that they are in the expected format. It will raise a MlflowException if
    the result is not in the expected format.

    :param custom_metric_tuple: Containing a user provided function and its index in the
                                ``custom_metrics`` parameter of ``mlflow.evaluate``
    :param eval_df: A Pandas dataframe object containing a prediction and a target column.
    :param builtin_metrics: A dictionary of metrics produced by the default evaluator.
    :return: A scalar metric value.
    """
    exception_header = (
        f"Custom metric '{custom_metric_tuple.name}' at index {custom_metric_tuple.index}"
        " in the `custom_metrics` parameter"
    )

    metric = custom_metric_tuple.function(eval_df, builtin_metrics)

    if metric is None:
        raise MlflowException(f"{exception_header} returned None.")

    if not _is_numeric(metric):
        raise MlflowException(f"{exception_header} did not return a scalar numeric value.")

    return metric


def _is_valid_artifacts(artifacts):
    return isinstance(artifacts, dict) and all(isinstance(k, str) for k in artifacts.keys())


def _evaluate_custom_artifacts(custom_artifact_tuple, eval_df, builtin_metrics):
    """
    This function calls the `custom_artifact` function and performs validations on the returned
    result to ensure that they are in the expected format. It will raise a MlflowException if
    the result is not in the expected format.

    :param custom_metric_tuple: Containing a user provided function and its index in the
                                ``custom_artifacts`` parameter of ``mlflow.evaluate``
    :param eval_df: A Pandas dataframe object containing a prediction and a target column.
    :param builtin_metrics: A dictionary of metrics produced by the default evaluator.
    :return: A dictionary of artifacts.
    """
    exception_header = (
        f"Custom artifact function '{custom_artifact_tuple.name}' "
        " at index {custom_artifact_tuple.index}"
        " in the `custom_artifacts` parameter"
    )
    artifacts = custom_artifact_tuple.function(
        eval_df, builtin_metrics, custom_artifact_tuple.artifacts_dir
    )

    if artifacts is None:
        raise MlflowException(f"{exception_header} returned None.")

    if not _is_valid_artifacts(artifacts):
        raise MlflowException(
            f"{exception_header} did not return artifacts as a dictionary of string artifact "
            "names with their corresponding objects."
        )

    return artifacts


def _compute_df_mode_or_mean(df):
    """
    Compute mean (for continuous columns) and compute mode (for other columns) for the
    input dataframe, return a dict, key is column name, value is the corresponding mode or
    mean value, this function calls `_is_continuous` to determine whether the
    column is continuous column.
    """
    continuous_cols = [c for c in df.columns if _is_continuous(df[c])]
    df_cont = df[continuous_cols]
    df_non_cont = df.drop(continuous_cols, axis=1)

    means = {} if df_cont.empty else df_cont.mean().to_dict()
    modes = {} if df_non_cont.empty else df_non_cont.mode().loc[0].to_dict()
    return {**means, **modes}


_SUPPORTED_SHAP_ALGORITHMS = ("exact", "permutation", "partition", "kernel")


def _shap_predict_fn(x, predict_fn, feature_names):
    return predict_fn(_get_dataframe_with_renamed_columns(x, feature_names))


# pylint: disable=attribute-defined-outside-init
class DefaultEvaluator(ModelEvaluator):
    # pylint: disable=unused-argument
    def can_evaluate(self, *, model_type, evaluator_config, **kwargs):
        return model_type in ["classifier", "regressor"]

    def _log_metrics(self):
        """
        Helper method to log metrics into specified run.
        """
        timestamp = get_current_time_millis()
        self.client.log_batch(
            self.run_id,
            metrics=[
                Metric(
                    key=key,
                    value=value,
                    timestamp=timestamp,
                    step=0,
                )
                for key, value in self.metrics.items()
            ],
        )

    def _log_image_artifact(
        self,
        do_plot,
        artifact_name,
    ):
        from matplotlib import pyplot

        prefix = self.evaluator_config.get("metric_prefix", "")
        artifact_file_name = f"{prefix}{artifact_name}.png"
        artifact_file_local_path = self.temp_dir.path(artifact_file_name)

        try:
            pyplot.clf()
            do_plot()
            pyplot.savefig(artifact_file_local_path, bbox_inches="tight")
        finally:
            pyplot.close(pyplot.gcf())

        mlflow.log_artifact(artifact_file_local_path)
        artifact = ImageEvaluationArtifact(uri=mlflow.get_artifact_uri(artifact_file_name))
        artifact._load(artifact_file_local_path)
        self.artifacts[artifact_name] = artifact

    def _log_pandas_df_artifact(self, pandas_df, artifact_name):
        artifact_file_name = f"{artifact_name}.csv"
        artifact_file_local_path = self.temp_dir.path(artifact_file_name)
        pandas_df.to_csv(artifact_file_local_path, index=False)
        mlflow.log_artifact(artifact_file_local_path)
        artifact = CsvEvaluationArtifact(
            uri=mlflow.get_artifact_uri(artifact_file_name),
            content=pandas_df,
        )
        artifact._load(artifact_file_local_path)
        self.artifacts[artifact_name] = artifact

    def _log_model_explainability(self):
        if not self.evaluator_config.get("log_model_explainability", True):
            return

        if self.is_model_server and not self.evaluator_config.get(
            "log_model_explainability", False
        ):
            _logger.warning(
                "Skipping model explainability because a model server is used for environment "
                "restoration."
            )
            return

        if self.model_loader_module == "mlflow.spark":
            # TODO: Shap explainer need to manipulate on each feature values,
            #  but spark model input dataframe contains Vector type feature column
            #  which shap explainer does not support.
            #  To support this, we need expand the Vector type feature column into
            #  multiple scalar feature columns and pass it to shap explainer.
            _logger.warning(
                "Logging model explainability insights is not currently supported for PySpark "
                "models."
            )
            return

        if not (np.issubdtype(self.y.dtype, np.number) or self.y.dtype == np.bool_):
            # Note: python bool type inherits number type but np.bool_ does not inherit np.number.
            _logger.warning(
                "Skip logging model explainability insights because it requires all label "
                "values to be numeric or boolean."
            )
            return

        algorithm = self.evaluator_config.get("explainability_algorithm", None)
        if algorithm is not None and algorithm not in _SUPPORTED_SHAP_ALGORITHMS:
            raise MlflowException(
                message=f"Specified explainer algorithm {algorithm} is unsupported. Currently only "
                f"support {','.join(_SUPPORTED_SHAP_ALGORITHMS)} algorithms.",
                error_code=INVALID_PARAMETER_VALUE,
            )

        if algorithm != "kernel":
            feature_dtypes = list(self.X.get_original().dtypes)
            for feature_dtype in feature_dtypes:
                if not np.issubdtype(feature_dtype, np.number):
                    _logger.warning(
                        "Skip logging model explainability insights because the shap explainer "
                        f"{algorithm} requires all feature values to be numeric, and each feature "
                        "column must only contain scalar values."
                    )
                    return

        try:
            import shap
            from matplotlib import pyplot
        except ImportError:
            _logger.warning(
                "SHAP or matplotlib package is not installed, so model explainability insights "
                "will not be logged."
            )
            return

        if Version(shap.__version__) < Version("0.40"):
            _logger.warning(
                "Shap package version is lower than 0.40, Skip log model explainability."
            )
            return

        is_multinomial_classifier = self.model_type == "classifier" and self.num_classes > 2

        sample_rows = self.evaluator_config.get(
            "explainability_nsamples", _DEFAULT_SAMPLE_ROWS_FOR_SHAP
        )

        X_df = self.X.copy_to_avoid_mutation()

        sampled_X = shap.sample(X_df, sample_rows, random_state=0)

        mode_or_mean_dict = _compute_df_mode_or_mean(X_df)
        sampled_X = sampled_X.fillna(mode_or_mean_dict)

        # shap explainer might call provided `predict_fn` with a `numpy.ndarray` type
        # argument, this might break some model inference, so convert the argument into
        # a pandas dataframe.
        # The `shap_predict_fn` calls model's predict function, we need to restore the input
        # dataframe with original column names, because some model prediction routine uses
        # the column name.

        shap_predict_fn = functools.partial(
            _shap_predict_fn, predict_fn=self.predict_fn, feature_names=self.feature_names
        )

        try:
            if algorithm:
                if algorithm == "kernel":
                    # We need to lazily import shap, so lazily import `_PatchedKernelExplainer`
                    from ._shap_patch import _PatchedKernelExplainer

                    kernel_link = self.evaluator_config.get(
                        "explainability_kernel_link", "identity"
                    )
                    if kernel_link not in ["identity", "logit"]:
                        raise ValueError(
                            "explainability_kernel_link config can only be set to 'identity' or "
                            f"'logit', but got '{kernel_link}'."
                        )
                    background_X = shap.sample(X_df, sample_rows, random_state=3)
                    background_X = background_X.fillna(mode_or_mean_dict)

                    explainer = _PatchedKernelExplainer(
                        shap_predict_fn, background_X, link=kernel_link
                    )
                else:
                    explainer = shap.Explainer(
                        shap_predict_fn,
                        sampled_X,
                        feature_names=self.feature_names,
                        algorithm=algorithm,
                    )
            else:
                if (
                    self.raw_model
                    and not is_multinomial_classifier
                    and not isinstance(self.raw_model, sk_Pipeline)
                ):
                    # For mulitnomial classifier, shap.Explainer may choose Tree/Linear explainer
                    # for raw model, this case shap plot doesn't support it well, so exclude the
                    # multinomial_classifier case here.
                    explainer = shap.Explainer(
                        self.raw_model, sampled_X, feature_names=self.feature_names
                    )
                else:
                    # fallback to default explainer
                    explainer = shap.Explainer(
                        shap_predict_fn, sampled_X, feature_names=self.feature_names
                    )

            _logger.info(f"Shap explainer {explainer.__class__.__name__} is used.")

            if algorithm == "kernel":
                shap_values = shap.Explanation(
                    explainer.shap_values(sampled_X), feature_names=self.feature_names
                )
            else:
                shap_values = explainer(sampled_X)
        except Exception as e:
            # Shap evaluation might fail on some edge cases, e.g., unsupported input data values
            # or unsupported model on specific shap explainer. Catch exception to prevent it
            # breaking the whole `evaluate` function.

            if not self.evaluator_config.get("ignore_exceptions", True):
                raise e

            _logger.warning(
                f"Shap evaluation failed. Reason: {repr(e)}. "
                "Set logging level to DEBUG to see the full traceback."
            )
            _logger.debug("", exc_info=True)
            return
        try:
            mlflow.shap.log_explainer(explainer, artifact_path="explainer")
        except Exception as e:
            # TODO: The explainer saver is buggy, if `get_underlying_model_flavor` return "unknown",
            #   then fallback to shap explainer saver, and shap explainer will call `model.save`
            #   for sklearn model, there is no `.save` method, so error will happen.
            _logger.warning(
                f"Logging explainer failed. Reason: {repr(e)}. "
                "Set logging level to DEBUG to see the full traceback."
            )
            _logger.debug("", exc_info=True)

        def _adjust_color_bar():
            pyplot.gcf().axes[-1].set_aspect("auto")
            pyplot.gcf().axes[-1].set_box_aspect(50)

        def _adjust_axis_tick():
            pyplot.xticks(fontsize=10)
            pyplot.yticks(fontsize=10)

        def plot_beeswarm():
            shap.plots.beeswarm(shap_values, show=False, color_bar=True)
            _adjust_color_bar()
            _adjust_axis_tick()

        self._log_image_artifact(
            plot_beeswarm,
            "shap_beeswarm_plot",
        )

        def plot_summary():
            shap.summary_plot(shap_values, show=False, color_bar=True)
            _adjust_color_bar()
            _adjust_axis_tick()

        self._log_image_artifact(
            plot_summary,
            "shap_summary_plot",
        )

        def plot_feature_importance():
            shap.plots.bar(shap_values, show=False)
            _adjust_axis_tick()

        self._log_image_artifact(
            plot_feature_importance,
            "shap_feature_importance_plot",
        )

    def _evaluate_sklearn_model_score_if_scorable(self):
        if self.model_loader_module == "mlflow.sklearn" and self.raw_model is not None:
            try:
                score = self.raw_model.score(
                    self.X.copy_to_avoid_mutation(), self.y, sample_weight=self.sample_weights
                )
                self.metrics["score"] = score
            except Exception as e:
                _logger.warning(
                    f"Computing sklearn model score failed: {repr(e)}. Set logging level to "
                    "DEBUG to see the full traceback."
                )
                _logger.debug("", exc_info=True)

    def _compute_roc_and_pr_curve(self):
        if self.y_probs is not None:
            self.roc_curve = _gen_classifier_curve(
                is_binomial=True,
                y=self.y,
                y_probs=self.y_prob,
                labels=self.label_list,
                pos_label=self.pos_label,
                curve_type="roc",
                sample_weights=self.sample_weights,
            )

            self.metrics["roc_auc"] = self.roc_curve.auc
            self.pr_curve = _gen_classifier_curve(
                is_binomial=True,
                y=self.y,
                y_probs=self.y_prob,
                labels=self.label_list,
                pos_label=self.pos_label,
                curve_type="pr",
                sample_weights=self.sample_weights,
            )

            self.metrics["precision_recall_auc"] = self.pr_curve.auc

    def _log_multiclass_classifier_artifacts(self):
        per_class_metrics_collection_df = _get_classifier_per_class_metrics_collection_df(
            self.y,
            self.y_pred,
            labels=self.label_list,
            sample_weights=self.sample_weights,
        )

        log_roc_pr_curve = False
        if self.y_probs is not None:
            max_classes_for_multiclass_roc_pr = self.evaluator_config.get(
                "max_classes_for_multiclass_roc_pr", 10
            )
            if self.num_classes <= max_classes_for_multiclass_roc_pr:
                log_roc_pr_curve = True
            else:
                _logger.warning(
                    f"The classifier num_classes > {max_classes_for_multiclass_roc_pr}, skip "
                    f"logging ROC curve and Precision-Recall curve. You can add evaluator config "
                    f"'max_classes_for_multiclass_roc_pr' to increase the threshold."
                )

        if log_roc_pr_curve:
            roc_curve = _gen_classifier_curve(
                is_binomial=False,
                y=self.y,
                y_probs=self.y_probs,
                labels=self.label_list,
                pos_label=self.pos_label,
                curve_type="roc",
                sample_weights=self.sample_weights,
            )

            def plot_roc_curve():
                roc_curve.plot_fn(**roc_curve.plot_fn_args)

            self._log_image_artifact(plot_roc_curve, "roc_curve_plot")
            per_class_metrics_collection_df["roc_auc"] = roc_curve.auc

            pr_curve = _gen_classifier_curve(
                is_binomial=False,
                y=self.y,
                y_probs=self.y_probs,
                labels=self.label_list,
                pos_label=self.pos_label,
                curve_type="pr",
                sample_weights=self.sample_weights,
            )

            def plot_pr_curve():
                pr_curve.plot_fn(**pr_curve.plot_fn_args)

            self._log_image_artifact(plot_pr_curve, "precision_recall_curve_plot")
            per_class_metrics_collection_df["precision_recall_auc"] = pr_curve.auc

        self._log_pandas_df_artifact(per_class_metrics_collection_df, "per_class_metrics")

    def _log_binary_classifier_artifacts(self):
        from mlflow.models.evaluation.lift_curve import plot_lift_curve

        if self.y_probs is not None:

            def plot_roc_curve():
                self.roc_curve.plot_fn(**self.roc_curve.plot_fn_args)

            self._log_image_artifact(plot_roc_curve, "roc_curve_plot")

            def plot_pr_curve():
                self.pr_curve.plot_fn(**self.pr_curve.plot_fn_args)

            self._log_image_artifact(plot_pr_curve, "precision_recall_curve_plot")

            self._log_image_artifact(
                lambda: plot_lift_curve(self.y, self.y_probs, pos_label=self.pos_label),
                "lift_curve_plot",
            )

    def _log_custom_metric_artifact(self, artifact_name, raw_artifact, custom_metric_tuple):
        """
        This function logs and returns a custom metric artifact. Two cases:
            - The provided artifact is a path to a file, the function will make a copy of it with
              a formatted name in a temporary directory and call mlflow.log_artifact.
            - Otherwise: will attempt to save the artifact to an temporary path with an inferred
              type. Then call mlflow.log_artifact.

        :param artifact_name: the name of the artifact
        :param raw_artifact:  the object representing the artifact
        :param custom_metric_tuple: an instance of the _CustomMetric namedtuple
        :return: EvaluationArtifact
        """

        exception_and_warning_header = (
            f"Custom metric function '{custom_metric_tuple.name}' at index "
            f"{custom_metric_tuple.index} in the `custom_metrics` parameter"
        )

        inferred_from_path, inferred_type, inferred_ext = _infer_artifact_type_and_ext(
            artifact_name, raw_artifact, custom_metric_tuple
        )
        artifact_file_local_path = self.temp_dir.path(artifact_name + inferred_ext)

        if pathlib.Path(artifact_file_local_path).exists():
            raise MlflowException(
                f"{exception_and_warning_header} produced an artifact '{artifact_name}' that "
                "cannot be logged because there already exists an artifact with the same name."
            )

        # ParquetEvaluationArtifact isn't explicitly stated here because such artifacts can only
        # be supplied through file. Which is handled by the first if clause. This is because
        # DataFrame objects default to be stored as CsvEvaluationArtifact.
        if inferred_from_path:
            shutil.copyfile(raw_artifact, artifact_file_local_path)
        elif inferred_type is JsonEvaluationArtifact:
            with open(artifact_file_local_path, "w") as f:
                if isinstance(raw_artifact, str):
                    f.write(raw_artifact)
                else:
                    json.dump(raw_artifact, f, cls=NumpyEncoder)
        elif inferred_type is CsvEvaluationArtifact:
            raw_artifact.to_csv(artifact_file_local_path, index=False)
        elif inferred_type is NumpyEvaluationArtifact:
            np.save(artifact_file_local_path, raw_artifact, allow_pickle=False)
        elif inferred_type is ImageEvaluationArtifact:
            raw_artifact.savefig(artifact_file_local_path)
        else:
            # storing as pickle
            try:
                with open(artifact_file_local_path, "wb") as f:
                    pickle.dump(raw_artifact, f)
                _logger.warning(
                    f"{exception_and_warning_header} produced an artifact '{artifact_name}'"
                    f" with type '{type(raw_artifact)}' that is logged as a pickle artifact."
                )
            except pickle.PickleError:
                raise MlflowException(
                    f"{exception_and_warning_header} produced an unsupported artifact "
                    f"'{artifact_name}' with type '{type(raw_artifact)}' that cannot be pickled. "
                    "Supported object types for artifacts are:\n"
                    "- A string uri representing the file path to the artifact. MLflow"
                    "  will infer the type of the artifact based on the file extension.\n"
                    "- A string representation of a JSON object. This will be saved as a "
                    ".json artifact.\n"
                    "- Pandas DataFrame. This will be saved as a .csv artifact."
                    "- Numpy array. This will be saved as a .npy artifact."
                    "- Matplotlib Figure. This will be saved as an .png image artifact."
                    "- Other objects will be attempted to be pickled with default protocol."
                )

        mlflow.log_artifact(artifact_file_local_path)
        artifact = inferred_type(uri=mlflow.get_artifact_uri(artifact_name + inferred_ext))
        artifact._load(artifact_file_local_path)
        return artifact

    def _evaluate_custom_metrics_and_log_produced_artifacts(self, log_to_mlflow_tracking=True):
        if not self.custom_metrics and not self.custom_artifacts:
            return
        builtin_metrics = copy.deepcopy(self.metrics)
        eval_df = pd.DataFrame({"prediction": copy.deepcopy(self.y_pred), "target": self.y})
        for index, custom_metric in enumerate(self.custom_metrics or []):
            # deepcopying eval_df and builtin_metrics for each custom metric function call,
            # in case the user modifies them inside their function(s).
            custom_metric_tuple = _CustomMetric(
                function=custom_metric.eval_fn,
                index=index,
                name=custom_metric.name,
            )
            metric_result = _evaluate_custom_metric(
                custom_metric_tuple,
                eval_df.copy(),
                copy.deepcopy(builtin_metrics),
            )
            self.metrics.update({custom_metric.name: metric_result})

        for index, custom_artifact in enumerate(self.custom_artifacts or []):
            with tempfile.TemporaryDirectory() as artifacts_dir:
                # deepcopying eval_df and builtin_metrics for each custom metric function call,
                # in case the user modifies them inside their function(s).
                custom_artifact_tuple = _CustomArtifact(
                    function=custom_artifact,
                    index=index,
                    name=getattr(custom_artifact, "__name__", repr(custom_artifact)),
                    artifacts_dir=artifacts_dir,
                )
                artifact_results = _evaluate_custom_artifacts(
                    custom_artifact_tuple,
                    eval_df.copy(),
                    copy.deepcopy(builtin_metrics),
                )
                if artifact_results is not None and log_to_mlflow_tracking:
                    for artifact_name, raw_artifact in artifact_results.items():
                        self.artifacts[artifact_name] = self._log_custom_metric_artifact(
                            artifact_name,
                            raw_artifact,
                            custom_artifact_tuple,
                        )

    def _log_confusion_matrix(self):
        """
        Helper method for logging confusion matrix
        """
        # normalize the confusion matrix, keep consistent with sklearn autologging.
        confusion_matrix = sk_metrics.confusion_matrix(
            self.y,
            self.y_pred,
            labels=self.label_list,
            normalize="true",
            sample_weight=self.sample_weights,
        )

        def plot_confusion_matrix():
            import matplotlib
            import matplotlib.pyplot as plt

            with matplotlib.rc_context(
                {
                    "font.size": min(8, math.ceil(50.0 / self.num_classes)),
                    "axes.labelsize": 8,
                }
            ):
                _, ax = plt.subplots(1, 1, figsize=(6.0, 4.0), dpi=175)
                disp = sk_metrics.ConfusionMatrixDisplay(
                    confusion_matrix=confusion_matrix,
                    display_labels=self.label_list,
                ).plot(cmap="Blues", ax=ax)
                disp.ax_.set_title("Normalized confusion matrix")

        if hasattr(sk_metrics, "ConfusionMatrixDisplay"):
            self._log_image_artifact(
                plot_confusion_matrix,
                "confusion_matrix",
            )
        return

    def _generate_model_predictions(self):
        """
        Helper method for generating model predictions
        """
        if self.model_type == "classifier":
            self.label_list = np.unique(self.y)
            self.num_classes = len(self.label_list)

            self.y_pred = self.predict_fn(self.X.copy_to_avoid_mutation())
            self.is_binomial = self.num_classes <= 2

            if self.is_binomial:
                if self.pos_label in self.label_list:
                    self.label_list = np.delete(
                        self.label_list, np.where(self.label_list == self.pos_label)
                    )
                    self.label_list = np.append(self.label_list, self.pos_label)
                elif self.pos_label is None:
                    self.pos_label = self.label_list[-1]
                _logger.info(
                    "The evaluation dataset is inferred as binary dataset, positive label is "
                    f"{self.label_list[1]}, negative label is {self.label_list[0]}."
                )
            else:
                _logger.info(
                    "The evaluation dataset is inferred as multiclass dataset, number of classes "
                    f"is inferred as {self.num_classes}"
                )

            if self.predict_proba_fn is not None:
                self.y_probs = self.predict_proba_fn(self.X.copy_to_avoid_mutation())
                if self.is_binomial:
                    self.y_prob = self.y_probs[:, 1]
                else:
                    self.y_prob = None
            else:
                self.y_probs = None
                self.y_prob = None
        elif self.model_type == "regressor":
            self.y_pred = self.model.predict(self.X.copy_to_avoid_mutation())

    def _compute_builtin_metrics(self):
        """
        Helper method for computing builtin metrics
        """
        self._evaluate_sklearn_model_score_if_scorable()
        if self.model_type == "classifier":
            if self.is_binomial:
                self.metrics.update(
                    _get_binary_classifier_metrics(
                        y_true=self.y,
                        y_pred=self.y_pred,
                        y_proba=self.y_probs,
                        labels=self.label_list,
                        pos_label=self.pos_label,
                        sample_weights=self.sample_weights,
                    )
                )
                self._compute_roc_and_pr_curve()
            else:
                average = self.evaluator_config.get("average", "weighted")
                self.metrics.update(
                    _get_multiclass_classifier_metrics(
                        y_true=self.y,
                        y_pred=self.y_pred,
                        y_proba=self.y_probs,
                        labels=self.label_list,
                        average=average,
                        sample_weights=self.sample_weights,
                    )
                )
        elif self.model_type == "regressor":
            self.metrics.update(_get_regressor_metrics(self.y, self.y_pred, self.sample_weights))

    def _log_metrics_and_artifacts(self):
        """
        Helper method for generating artifacts, logging metrics and artifacts.
        """
        if self.model_type == "classifier":
            if self.is_binomial:
                self._log_binary_classifier_artifacts()
            else:
                self._log_multiclass_classifier_artifacts()
            self._log_confusion_matrix()
        self._log_metrics()
        self._log_model_explainability()

    def _evaluate(
        self,
        model: "mlflow.pyfunc.PyFuncModel",
        is_baseline_model=False,
        **kwargs,
    ):
        import matplotlib

        with TempDir() as temp_dir, matplotlib.rc_context(_matplotlib_config):
            self.client = MlflowClient()

            self.temp_dir = temp_dir
            self.model = model
            self.is_baseline_model = is_baseline_model

            self.is_model_server = isinstance(model, _ServedPyFuncModel)

            model_loader_module, raw_model = _extract_raw_model(model)
            predict_fn, predict_proba_fn = _extract_predict_fn(model, raw_model)

            self.model_loader_module = model_loader_module
            self.raw_model = raw_model
            self.predict_fn = predict_fn
            self.predict_proba_fn = predict_proba_fn

            self.metrics = {}
            self.baseline_metrics = {}
            self.artifacts = {}

            if self.model_type not in ["classifier", "regressor"]:
                raise MlflowException(
                    message=f"Unsupported model type {self.model_type}",
                    erorr_code=INVALID_PARAMETER_VALUE,
                )
            with mlflow.utils.autologging_utils.disable_autologging():
                self._generate_model_predictions()
                self._compute_builtin_metrics()
                self._evaluate_custom_metrics_and_log_produced_artifacts(
                    log_to_mlflow_tracking=not is_baseline_model
                )
                metric_prefix = self.evaluator_config.get("metric_prefix")
                if metric_prefix is not None:
                    self.metrics = {
                        f"{metric_prefix}{metric_key}": metric_value
                        for metric_key, metric_value in self.metrics.items()
                    }
                if not is_baseline_model:
                    self._log_metrics_and_artifacts()
                return EvaluationResult(metrics=self.metrics, artifacts=self.artifacts)

    def evaluate(
        self,
        *,
        model: "mlflow.pyfunc.PyFuncModel",
        model_type,
        dataset,
        run_id,
        evaluator_config,
        custom_metrics=None,
        custom_artifacts=None,
        baseline_model=None,
        **kwargs,
    ):
        self.dataset = dataset
        self.run_id = run_id
        self.model_type = model_type
        self.evaluator_config = evaluator_config
        self.dataset_name = dataset.name
        self.feature_names = dataset.feature_names
        self.custom_metrics = custom_metrics
        self.custom_artifacts = custom_artifacts
        self.y = dataset.labels_data
        self.pos_label = self.evaluator_config.get("pos_label")
        self.sample_weights = self.evaluator_config.get("sample_weights")

        inferred_model_type = _infer_model_type_by_labels(self.y)

        if inferred_model_type is not None and model_type != inferred_model_type:
            _logger.warning(
                f"According to the evaluation dataset label values, the model type looks like "
                f"{inferred_model_type}, but you specified model type {model_type}. Please "
                f"verify that you set the `model_type` and `dataset` arguments correctly."
            )

        if evaluator_config.get("_disable_candidate_model", False):
            evaluation_result = EvaluationResult(metrics={}, artifacts={})
        else:
            if baseline_model:
                _logger.info("Evaluating candidate model:")
            evaluation_result = self._evaluate(model, is_baseline_model=False)

        if not baseline_model:
            return evaluation_result

        _logger.info("Evaluating baseline model:")
        baseline_evaluation_result = self._evaluate(baseline_model, is_baseline_model=True)

        return EvaluationResult(
            metrics=evaluation_result.metrics,
            artifacts=evaluation_result.artifacts,
            baseline_model_metrics=baseline_evaluation_result.metrics,
        )

    @property
    def X(self) -> pd.DataFrame:
        """
        The features (`X`) portion of the dataset, guarded against accidental mutations.
        """
        return DefaultEvaluator._MutationGuardedData(
            _get_dataframe_with_renamed_columns(self.dataset.features_data, self.feature_names)
        )

    class _MutationGuardedData:
        """
        Wrapper around a data object that requires explicit API calls to obtain either a copy
        of the data object, or, in cases where the caller can guaranteed that the object will not
        be mutated, the original data object.
        """

        def __init__(self, data):
            """
            :param data: A data object, such as a Pandas DataFrame, numPy array, or list.
            """
            self._data = data

        def copy_to_avoid_mutation(self):
            """
            Obtain a copy of the data. This method should be called every time the data needs
            to be used in a context where it may be subsequently mutated, guarding against
            accidental reuse after mutation.

            :return: A copy of the data object.
            """
            if isinstance(self._data, pd.DataFrame):
                return self._data.copy(deep=True)
            else:
                return copy.deepcopy(self._data)

        def get_original(self):
            """
            Obtain the original data object. This method should only be called if the caller
            can guarantee that it will not mutate the data during subsequent operations.

            :return: The original data object.
            """
            return self._data
