import datetime
import logging
import operator
import os
import sys
import warnings
from collections import namedtuple
from pathlib import Path
from typing import Any

import mlflow
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
from mlflow.recipes.cards import BaseCard
from mlflow.recipes.step import BaseStep, StepClass
from mlflow.recipes.steps.train import TrainStep
from mlflow.recipes.utils.execution import get_step_output_path
from mlflow.recipes.utils.metrics import (
    _get_builtin_metrics,
    _get_custom_metrics,
    _get_extended_task,
    _get_model_type_from_template,
    _get_primary_metric,
    _load_custom_metrics,
    transform_multiclass_metric,
)
from mlflow.recipes.utils.step import get_merged_eval_metrics, validate_classification_config
from mlflow.recipes.utils.tracking import (
    TrackingConfig,
    apply_recipe_tracking_config,
    get_recipe_tracking_config,
    get_run_tags_env_vars,
)
from mlflow.tracking.fluent import _get_experiment_id, _set_experiment_primary_metric
from mlflow.utils.databricks_utils import get_databricks_env_vars, get_databricks_run_url
from mlflow.utils.string_utils import strip_prefix

_logger = logging.getLogger(__name__)


_FEATURE_IMPORTANCE_PLOT_FILE = "feature_importance.png"
_VALIDATION_METRIC_PREFIX = "val_"

MetricValidationResult = namedtuple(
    "MetricValidationResult", ["metric", "greater_is_better", "value", "threshold", "validated"]
)


class EvaluateStep(BaseStep):
    def __init__(self, step_config: dict[str, Any], recipe_root: str) -> None:
        super().__init__(step_config, recipe_root)
        self.tracking_config = TrackingConfig.from_dict(self.step_config)

    def _validate_and_apply_step_config(self):
        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.model_validation_status = "UNKNOWN"
        self.primary_metric = _get_primary_metric(
            self.step_config.get("primary_metric"), self.extended_task
        )
        self.user_defined_custom_metrics = {
            metric.name: metric
            for metric in _get_custom_metrics(self.step_config, self.extended_task)
        }
        self.evaluation_metrics = {
            metric.name: metric for metric in _get_builtin_metrics(self.extended_task)
        }
        self.evaluation_metrics.update(self.user_defined_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,
            )

    def _validate_validation_criteria(self):
        """
        Validates validation criteria don't contain undefined metrics
        """
        val_metrics = {vc["metric"] for vc in self.step_config.get("validation_criteria", [])}
        if not val_metrics:
            return
        undefined_metrics = val_metrics.difference(self.evaluation_metrics.keys())
        if undefined_metrics:
            raise MlflowException(
                f"Validation criteria contain undefined metrics: {sorted(undefined_metrics)}",
                error_code=INVALID_PARAMETER_VALUE,
            )

    def _check_validation_criteria(self, metrics, validation_criteria):
        """
        return a list of `MetricValidationResult` tuple instances.
        """
        summary = []
        for val_criterion in validation_criteria:
            metric_name = val_criterion["metric"]
            metric_val = metrics.get(metric_name)
            if metric_val is None:
                raise MlflowException(
                    f"The metric {metric_name} is defined in the recipe's validation criteria"
                    " but was not returned from mlflow evaluation.",
                    error_code=INVALID_PARAMETER_VALUE,
                )
            greater_is_better = self.evaluation_metrics[metric_name].greater_is_better
            comp_func = operator.ge if greater_is_better else operator.le
            threshold = val_criterion["threshold"]
            validated = comp_func(metric_val, threshold)
            summary.append(
                MetricValidationResult(
                    metric=metric_name,
                    greater_is_better=greater_is_better,
                    value=metric_val,
                    threshold=threshold,
                    validated=validated,
                )
            )
        return summary

    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 pandas as pd

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

            self._validate_validation_criteria()

            test_df_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="split",
                relative_path="test.parquet",
            )
            test_df = pd.read_parquet(test_df_path)
            validate_classification_config(self.task, self.positive_class, test_df, self.target_col)

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

            run_id_path = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="train",
                relative_path="run_id",
            )
            run_id = Path(run_id_path).read_text()

            model_uri = get_step_output_path(
                recipe_root_path=self.recipe_root,
                step_name="train",
                relative_path=TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH,
            )

            apply_recipe_tracking_config(self.tracking_config)
            exp_id = _get_experiment_id()

            primary_metric_greater_is_better = self.evaluation_metrics[
                self.primary_metric
            ].greater_is_better

            _set_experiment_primary_metric(
                exp_id, f"test_{self.primary_metric}", primary_metric_greater_is_better
            )

            with mlflow.start_run(run_id=run_id):
                eval_metrics = {}
                for dataset_name, dataset, evaluator_config in (
                    (
                        "validation",
                        validation_df,
                        {
                            "explainability_algorithm": "kernel",
                            "explainability_nsamples": 10,
                            "metric_prefix": _VALIDATION_METRIC_PREFIX,
                        },
                    ),
                    (
                        "test",
                        test_df,
                        {
                            "log_model_explainability": False,
                            "metric_prefix": "test_",
                        },
                    ),
                ):
                    if self.extended_task == "classification/binary":
                        evaluator_config["pos_label"] = self.positive_class
                    eval_result = mlflow.evaluate(
                        model=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=evaluator_config,
                    )
                    eval_result.save(os.path.join(output_directory, f"eval_{dataset_name}"))
                    eval_metrics[dataset_name] = {
                        transform_multiclass_metric(
                            strip_prefix(k, evaluator_config["metric_prefix"]), self.extended_task
                        ): v
                        for k, v in eval_result.metrics.items()
                    }

                validation_results = self._validate_model(eval_metrics, output_directory)

            card = self._build_profiles_and_card(
                run_id, model_uri, eval_metrics, validation_results, output_directory
            )
            card.save_as_html(output_directory)
            self._log_step_card(run_id, self.name)
            return card
        finally:
            warnings.warn = original_warn

    def _validate_model(self, eval_metrics, output_directory):
        validation_criteria = self.step_config.get("validation_criteria")
        validation_results = None
        if validation_criteria:
            validation_results = self._check_validation_criteria(
                eval_metrics["test"], validation_criteria
            )
            self.model_validation_status = (
                "VALIDATED" if all(cr.validated for cr in validation_results) else "REJECTED"
            )
        else:
            self.model_validation_status = "UNKNOWN"
        Path(output_directory, "model_validation_status").write_text(self.model_validation_status)
        return validation_results

    def _build_profiles_and_card(
        self, run_id, model_uri, eval_metrics, validation_results, output_directory
    ):
        """
        Constructs data profiles of predictions and errors and a step card instance corresponding
        to the current evaluate step state.

        Args:
            run_id: The ID of the MLflow Run to which to log model evaluation results.
            model_uri: The URI of the model being evaluated.
            eval_metrics: The evaluation result keyed by dataset name from `mlflow.evaluate`.
            validation_results: A list of `MetricValidationResult` instances.
            output_directory: Output directory used by the evaluate step.
        """
        import pandas as pd

        # Build card
        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],
            )
            .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
            )
        )

        card.add_tab(
            "Model Performance (Test)",
            "<h3 class='section-title'>Summary Metrics</h3>"
            "<b>NOTE</b>: Use evaluation metrics over test dataset with care. "
            "Fine-tuning model over the test dataset is not advised."
            "{{ METRICS }} ",
        ).add_html("METRICS", metric_table_html)

        # Tab 1: model validation results, if exists.
        if validation_results is not None:

            def get_icon(validated):
                return (
                    # check mark button emoji
                    "\u2705"
                    if validated
                    # cross mark emoji
                    else "\u274c"
                )

            result_df = pd.DataFrame(validation_results).assign(
                validated=lambda df: df["validated"].map(get_icon)
            )

            criteria_html = BaseCard.render_table(
                result_df.style.format({"value": "{:.6g}", "threshold": "{:.6g}"})
            )
            card.add_tab("Model Validation", "{{ METRIC_VALIDATION_RESULTS }}").add_html(
                "METRIC_VALIDATION_RESULTS",
                "<h3 class='section-title'>Model Validation Results (Test Dataset)</h3> "
                + criteria_html,
            )

        # Tab 2: Classifier plots.
        if self.recipe == "classification/v1":
            classifiers_plot_tab = card.add_tab(
                "Model Performance Plots",
                "{{ CONFUSION_MATRIX }} {{CONFUSION_MATRIX_PLOT}}"
                + "{{ LIFT_CURVE }} {{LIFT_CURVE_PLOT}}"
                + "{{ PR_CURVE }} {{PR_CURVE_PLOT}}"
                + "{{ ROC_CURVE }} {{ROC_CURVE_PLOT}}",
            )
            confusion_matrix_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}confusion_matrix.png",
            )
            if os.path.exists(confusion_matrix_path):
                classifiers_plot_tab.add_html(
                    "CONFUSION_MATRIX",
                    '<h3 class="section-title">Confusion Matrix Plot</h3>',
                )
                classifiers_plot_tab.add_image(
                    "CONFUSION_MATRIX_PLOT", confusion_matrix_path, width=400
                )

            lift_curve_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}lift_curve_plot.png",
            )
            if os.path.exists(lift_curve_path):
                classifiers_plot_tab.add_html(
                    "LIFT_CURVE",
                    '<h3 class="section-title">Lift Curve Plot</h3>',
                )
                classifiers_plot_tab.add_image("LIFT_CURVE_PLOT", lift_curve_path, width=400)

            pr_curve_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}precision_recall_curve_plot.png",
            )
            if os.path.exists(pr_curve_path):
                classifiers_plot_tab.add_html(
                    "PR_CURVE",
                    '<h3 class="section-title">Precision Recall Curve Plot</h3>',
                )
                classifiers_plot_tab.add_image("PR_CURVE_PLOT", pr_curve_path, width=400)

            roc_curve_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}roc_curve_plot.png",
            )
            if os.path.exists(roc_curve_path):
                classifiers_plot_tab.add_html(
                    "ROC_CURVE",
                    '<h3 class="section-title">ROC Curve Plot</h3>',
                )
                classifiers_plot_tab.add_image("ROC_CURVE_PLOT", roc_curve_path, width=400)

        # Tab 3: SHAP plots.
        def _add_shap_plots(card):
            """Contingent on shap being installed."""
            shap_plot_tab = card.add_tab(
                "Feature Importance",
                '<h3 class="section-title">Feature Importance on Validation Dataset</h3>'
                '<h3 class="section-title">SHAP Bar Plot</h3>{{SHAP_BAR_PLOT}}'
                '<h3 class="section-title">SHAP Beeswarm Plot</h3>{{SHAP_BEESWARM_PLOT}}',
            )

            shap_bar_plot_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}shap_feature_importance_plot.png",
            )
            shap_beeswarm_plot_path = os.path.join(
                output_directory,
                "eval_validation/artifacts",
                f"{_VALIDATION_METRIC_PREFIX}shap_beeswarm_plot.png",
            )
            shap_plot_tab.add_image("SHAP_BAR_PLOT", shap_bar_plot_path, width=800)
            shap_plot_tab.add_image("SHAP_BEESWARM_PLOT", shap_beeswarm_plot_path, width=800)

        try:
            import shap  # noqa: F401
            from matplotlib import pyplot  # noqa: F401

            _add_shap_plots(card)
        except ImportError:
            _logger.warning(
                "SHAP or matplotlib package is not installed, so shap plots will not be added."
            )

        # Tab 3: 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 4: Run summary.
        run_summary_card_tab = card.add_tab(
            "Run Summary",
            "{{ RUN_ID }} "
            + "{{ MODEL_URI }}"
            + "{{ VALIDATION_STATUS }}"
            + "{{ EXE_DURATION }}"
            + "{{ LAST_UPDATE_TIME }}",
        ).add_markdown(
            "VALIDATION_STATUS", f"**Validation status:** `{self.model_validation_status}`"
        )
        run_url = get_databricks_run_url(
            tracking_uri=mlflow.get_tracking_uri(),
            run_id=run_id,
        )
        model_uri = f"runs:/{run_id}/train/{TrainStep.MODEL_ARTIFACT_RELATIVE_PATH}"
        model_url = get_databricks_run_url(
            tracking_uri=mlflow.get_tracking_uri(),
            run_id=run_id,
            artifact_path=f"train/{TrainStep.MODEL_ARTIFACT_RELATIVE_PATH}",
        )

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

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

        return card

    @classmethod
    def from_recipe_config(cls, recipe_config, recipe_root):
        step_config = {}
        if recipe_config.get("steps", {}).get("evaluate", {}) is not None:
            step_config.update(recipe_config.get("steps", {}).get("evaluate", {}))
        step_config["target_col"] = recipe_config.get("target_col")
        if "positive_class" in recipe_config:
            step_config["positive_class"] = recipe_config.get("positive_class")
        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.update(
            get_recipe_tracking_config(
                recipe_root_path=recipe_root,
                recipe_config=recipe_config,
            ).to_dict()
        )
        return cls(step_config, recipe_root)

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

    @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 step_class(self):
        return StepClass.TRAINING