#  Copyright (c) 2021 Rikai Authors
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Custom Mlflow model logger to make sure models have the right logging for
Rikai SQL ML
"""

import os
import tempfile
import warnings
from typing import Any, Optional

CONF_MLFLOW_TRACKING_URI = "spark.rikai.sql.ml.registry.mlflow.tracking_uri"
CONF_MLFLOW_OUTPUT_SCHEMA = "rikai.output.schema"
CONF_MLFLOW_SPEC_VERSION = "rikai.spec.version"
CONF_MLFLOW_PRE_PROCESSING = "rikai.transforms.pre"
CONF_MLFLOW_POST_PROCESSING = "rikai.transforms.post"
CONF_MLFLOW_MODEL_FLAVOR = "rikai.model.flavor"
CONF_MLFLOW_MODEL_TYPE = "rikai.model.type"
CONF_MLFLOW_ARTIFACT_PATH = "rikai.model.artifact_path"


class MlflowLogger:
    """
    An alternative model logger for use during training instead of the vanilla
    mlflow logger.

    """

    _CURRENT_MODEL_SPEC_VERSION = "1.0"

    def __init__(self, flavor: str):
        """
        Parameters
        ----------
        flavor: str
            The model flavor
        """
        self.spec_version = MlflowLogger._CURRENT_MODEL_SPEC_VERSION
        self.flavor = flavor

    def _log_tensorflow_model(
        self,
        model: Any,
        artifact_path: str,
        registered_model_name: str,
        **kwargs,
    ):
        import tensorflow as tf

        try:
            import mlflow
        except ImportError as e:
            raise ImportError(
                "Couldn't import mlflow. Please make sure to "
                "`pip install mlflow` explicitly or install "
                "the correct extras like `pip install rikai[mlflow]`"
            ) from e

        with tempfile.TemporaryDirectory() as tmp_dir:

            tf.saved_model.save(model, tmp_dir)

            mlflow.tensorflow.log_model(
                tf_saved_model_dir=tmp_dir,
                tf_meta_graph_tags=[tf.saved_model.SERVING],
                tf_signature_def_key=next(iter(model.signatures.keys())),
                artifact_path=artifact_path,
                registered_model_name=registered_model_name,
                **kwargs,
            )

    def log_model(
        self,
        model: Any,
        artifact_path: str,
        schema: Optional[str] = None,
        pre_processing: Optional[str] = None,
        post_processing: Optional[str] = None,
        registered_model_name: Optional[str] = None,
        customized_flavor: Optional[str] = None,
        model_type: Optional[str] = None,
        **kwargs,
    ):
        """Convenience function to log the model with tags needed by rikai.
        This should be called during training when the model is produced.

        Parameters
        ----------
        model: Any
            The model artifact object
        artifact_path: str
            The relative (to the run) artifact path
        schema: str
            Output schema (pyspark DataType)
        pre_processing: str, default None
            Full python module path of the pre-processing transforms
        post_processing: str, default None
            Full python module path of the post-processing transforms
        registered_model_name: str, default None
            Model name in the mlflow model registry
        kwargs: dict
            Passed to `mlflow.<flavor>.log_model`


        Examples
        --------

        .. code-block:: python

            import rikai.mlflow

            # Log PyTorch model
            with mlflow.start_run() as run:

                # Training loop
                # ...

                # Assume `model` is the trained model from the training loop
                rikai.mlflow.pytorch.log_model(model, "model",
                        schema=rikai.contrib.torch.transforms.fasterrcnn_resnet50_fpn.OUTPUT_SCHEMA,
                        pre_processing="rikai.contrib.torch.transforms.fasterrcnn_resnet50_fpn.pre_processing",
                        post_processing="rikai.contrib.torch.transforms.fasterrcnn_resnet50_fpn.post_processing",
                        registered_model_name="MyPytorchModel")


        For more details see `mlflow docs <https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html#mlflow.pytorch.log_model>`_.
        """  # noqa E501

        try:
            import mlflow
            from mlflow.tracking import MlflowClient
        except ImportError as e:
            raise ImportError(
                "Couldn't import mlflow. Please make sure to "
                "`pip install mlflow` explicitly or install "
                "the correct extras like `pip install rikai[mlflow]`"
            ) from e

        # no need to set the tracking uri here since this is intended to be
        # called inside the training loop within mlflow.start_run
        if self.flavor == "tensorflow":
            self._log_tensorflow_model(
                model,
                artifact_path,
                registered_model_name=registered_model_name,
                **kwargs,
            )
        else:
            getattr(mlflow, self.flavor).log_model(
                model,
                artifact_path,
                registered_model_name=registered_model_name,
                **kwargs,
            )

        tags = {
            CONF_MLFLOW_SPEC_VERSION: MlflowLogger._CURRENT_MODEL_SPEC_VERSION,
            CONF_MLFLOW_MODEL_FLAVOR: customized_flavor
            if customized_flavor
            else self.flavor,
            CONF_MLFLOW_MODEL_TYPE: model_type,
            CONF_MLFLOW_OUTPUT_SCHEMA: schema,
            CONF_MLFLOW_PRE_PROCESSING: pre_processing,
            CONF_MLFLOW_POST_PROCESSING: post_processing,
            CONF_MLFLOW_ARTIFACT_PATH: artifact_path,
        }
        for k in (
            CONF_MLFLOW_PRE_PROCESSING,
            CONF_MLFLOW_POST_PROCESSING,
            CONF_MLFLOW_MODEL_TYPE,
            CONF_MLFLOW_OUTPUT_SCHEMA,
        ):
            if not tags[k]:
                del tags[k]
                warnings.warn(
                    f"value of {k} is None or empty and "
                    "will not be populated to MLflow"
                )
        mlflow.set_tags(tags)
        if registered_model_name is not None:
            # if we're creating a model registry entry,
            # we also want to set the tags on the model version
            # and model to enable search
            c = MlflowClient()
            # mlflow log_model does not return the version (wtf)
            all_versions = c.get_latest_versions(
                registered_model_name, stages=["production", "staging", "None"]
            )
            current_version = None
            run_id = mlflow.active_run().info.run_id
            for v in all_versions:
                if v.run_id == run_id:
                    current_version = v
                    break
            if current_version is None:
                raise ValueError(
                    "No model version found matching runid: {}".format(run_id)
                )
            for key, value in tags.items():
                c.set_registered_model_tag(registered_model_name, key, value)
                c.set_model_version_tag(
                    registered_model_name, current_version.version, key, value
                )


KNOWN_FLAVORS = ["pytorch", "sklearn", "tensorflow"]
for flavor in KNOWN_FLAVORS:
    globals()[flavor] = MlflowLogger(flavor)
