#!/usr/bin/env python
# -*- coding: utf-8 -*--

# Copyright (c) 2022 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/

import json
import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import cloudpickle
import numpy as np
import pandas as pd
import requests
from ads.catalog.model import ModelCatalog
from ads.common import auth as authutil
from ads.common import logger, oci_client
from ads.common.model_export_util import ONNXTransformer
from ads.common.model_introspect import (
    TEST_STATUS,
    Introspectable,
    IntrospectionNotPassed,
    ModelIntrospect,
)
from ads.common.model_metadata import (
    Framework,
    ModelCustomMetadata,
    ModelProvenanceMetadata,
    ModelTaxonomyMetadata,
)
from ads.common.model_metadata_mixin import MetadataMixin
from ads.common.utils import DATA_SCHEMA_MAX_COL_NUM, get_files
from ads.config import (
    CONDA_BUCKET_NS,
    JOB_RUN_COMPARTMENT_OCID,
    JOB_RUN_OCID,
    NB_SESSION_COMPARTMENT_OCID,
    NB_SESSION_OCID,
    PROJECT_OCID,
)
from ads.feature_engineering.schema import Schema
from ads.model.artifact import ModelArtifact
from ads.model.common.utils import _extract_locals
from ads.model.deployment import (
    DEFAULT_POLL_INTERVAL,
    DEFAULT_WAIT_TIME,
    ModelDeployer,
    ModelDeployment,
    ModelDeploymentProperties,
)
from ads.model.deployment.common.utils import State as ModelDeploymentState
from ads.model.model_properties import ModelProperties
from ads.model.runtime.env_info import DEFAULT_CONDA_BUCKET_NAME
from ads.model.runtime.runtime_info import RuntimeInfo
from IPython.core.display import display
from enum import Enum, auto


_TRAINING_RESOURCE_ID = JOB_RUN_OCID or NB_SESSION_OCID
_COMPARTMENT_OCID = NB_SESSION_COMPARTMENT_OCID or JOB_RUN_COMPARTMENT_OCID

MODEL_DEPLOYMENT_INSTANCE_SHAPE = "VM.Standard2.1"
MODEL_DEPLOYMENT_INSTANCE_COUNT = 1
MODEL_DEPLOYMENT_BANDWIDTH_MBPS = 10


ONNX_DATA_TRANSFORMER = "onnx_data_transformer.json"


class NotActiveDeploymentError(Exception):
    def __init__(self, state: str):
        msg = (
            "To perform a prediction the deployed model needs to be in an active state. "
            f"The current state is: {state}."
        )
        super().__init__(msg)


class SerializeModelNotImplementedError(NotImplementedError):
    pass


class SerializeInputNotImplementedError(NotImplementedError):
    pass


class RuntimeInfoInconsistencyError(Exception):
    pass


class GenericModel(MetadataMixin, Introspectable):
    """Generic Model class which is the base class for all the frameworks including
    the unsupported frameworks.

    Attributes
    ----------
    algorithm: str
        The algorithm of the model.
    artifact_dir: str
        Artifact directory to store the files needed for deployment.
    auth: Dict
        Default authentication is set using the `ads.set_auth` API. To override the
        default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create
        an authentication signer to instantiate an IdentityClient object.
    ds_client: DataScienceClient
        The data science client used by model deployment.
    estimator: Callable
        Any model object generated by sklearn framework
    framework: str
        The framework of the model.
    hyperparameter: dict
        The hyperparameters of the estimator.
    metadata_custom: ModelCustomMetadata
        The model custom metadata.
    metadata_provenance: ModelProvenanceMetadata
        The model provenance metadata.
    metadata_taxonomy: ModelTaxonomyMetadata
        The model taxonomy metadata.
    model_artifact: ModelArtifact
        This is built by calling prepare.
    model_deployment: ModelDeployment
        A ModelDeployment instance.
    model_file_name: str
        Name of the serialized model.
    model_id: str
        The model ID.
    properties: ModelProperties
        ModelProperties object required to save and deploy model.
    runtime_info: RuntimeInfo
        A RuntimeInfo instance.
    schema_input: Schema
        Schema describes the structure of the input data.
    schema_output: Schema
        Schema describes the structure of the output data.
    serialize: bool
        Whether to serialize the model to pkl file by default. If False, you need to serialize the model manually,
        save it under artifact_dir and update the score.py manually.
    version: str
        The framework version of the model.

    Methods
    -------
    delete_deployment(...)
        Deletes the current model deployment.
    deploy(..., **kwargs)
        Deploys a model.
    from_model_artifact(uri, model_file_name, artifact_dir, ..., **kwargs)
        Loads model from the specified folder, or zip/tar archive.
    from_model_catalog(model_id, model_file_name, artifact_dir, ..., **kwargs)
        Loads model from model catalog.
    introspect(...)
        Runs model introspection.
    predict(data, ...)
        Returns prediction of input data run against the model deployment endpoint.
    prepare(..., **kwargs)
        Prepare and save the score.py, serialized model and runtime.yaml file.
    reload(...)
        Reloads the model artifact files: `score.py` and the `runtime.yaml`.
    save(..., **kwargs)
        Saves model artifacts to the model catalog.
    summary_status(...)
        Gets a summary table of the current status.
    verify(data, ...)
        Tests if deployment works in local environment.


    Examples
    --------
    >>> import tempfile
    >>> from ads.model.generic_model import GenericModel

    >>> class Toy:
    ...     def predict(self, x):
    ...         return x ** 2
    >>> estimator = Toy()

    >>> model = GenericModel(estimator=estimator, artifact_dir=tempfile.mkdtemp())
    >>> model.summary_status()
    >>> model.prepare(inference_conda_env="oci://service-conda-packs@id19sfcrra6z/service_pack/cpu/Data Exploration and Manipulation for CPU Python 3.7/3.0/dataexpl_p37_cpu_v3",
    ...              inference_python_version="3.7",
    ...              model_file_name="toy_model.pkl",
    ...              training_id=None,
    ...              force_overwrite=True
    ...            )
    >>> model.verify(2)
    >>> model.save()
    >>> model.deploy()
    >>> model.predict(2)
    >>> model.delete_deployment()
    """

    def __init__(
        self,
        estimator: Callable,
        artifact_dir: str,
        properties: Optional[ModelProperties] = None,
        auth: Optional[Dict] = None,
        serialize: bool = True,
        **kwargs: dict,
    ):
        """GenericModel Constructor.

        Parameters
        ----------
        estimator: (Callable).
            Trained model.
        artifact_dir: str
            Artifact directory to store the files needed for deployment.
        properties: (ModelProperties, optional). Defaults to None.
            ModelProperties object required to save and deploy model.
        auth :(Dict, optional). Defaults to None.
            The default authetication is set using `ads.set_auth` API. If you need to override the
            default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
            authentication signer and kwargs required to instantiate IdentityClient object.
        serialize: (bool, optional). Defaults to True.
            Whether to serialize the model to pkl file by default. If False, you need to serialize the model manually,
            save it under artifact_dir and update the score.py manually.
        """
        self.estimator = estimator
        self.auth = auth or authutil.default_signer()
        self.ds_client = oci_client.OCIClientFactory(**self.auth).data_science
        self.metadata_custom = ModelCustomMetadata()
        self.metadata_taxonomy = ModelTaxonomyMetadata()
        self.metadata_provenance = ModelProvenanceMetadata()
        self.schema_input = Schema()
        self.schema_output = Schema()
        self.model_file_name = None
        self.artifact_dir = os.path.abspath(os.path.expanduser(artifact_dir))
        self.model_artifact = None
        self.framework = None
        self.algorithm = None
        self.version = None
        self.hyperparameter = None
        self._introspect = ModelIntrospect(self)
        self.model_id = None
        self.model_deployment = None
        self.runtime_info = None
        self._as_onnx = kwargs.pop("as_onnx", False)

        if properties:
            self.properties = (
                properties
                if isinstance(properties, ModelProperties)
                else ModelProperties.from_dict(properties)
            )
        else:
            self.properties = ModelProperties().with_env()

        self._serialize = serialize
        self._summary_status = SummaryStatus()

    def serialize_model(
        self,
        as_onnx: bool = False,
        initial_types: List[Tuple] = None,
        force_overwrite: bool = False,
        X_sample: any = None,
    ):
        """
        Serialize and save model using ONNX or model specific method.

        Parameters
        ----------
        as_onnx: (boolean, optional)
            If set as True, convert into ONNX model.
        initial_types: (List[Tuple], optional)
            a python list. Each element is a tuple of a variable name and a data type.
        force_overwrite: (boolean, optional)
            If set as True, overwrite serialized model if exists.
        X_sample: (any, optional). Defaults to None.
            Contains model inputs such that model(X_sample) is a valid
            invocation of the model, used to valid model input type.

        Returns
        -------
        None
            Nothing
        """
        if self._serialize:
            with open(
                os.path.join(self.artifact_dir, self.model_file_name), "wb"
            ) as outfile:
                cloudpickle.dump(self.estimator, outfile)
        else:
            raise SerializeModelNotImplementedError(
                "`serialize_model` is not implemented."
            )

    def _handle_model_file_name(self, as_onnx: bool, model_file_name: str):
        if not model_file_name:
            if not self._serialize:
                raise NotImplementedError("`model_file_name` has to be provided.")
            else:
                model_file_name = "model.pkl"
        if as_onnx:
            assert model_file_name.endswith(
                ".onnx"
            ), "Wrong file extension. Expecting .onnx suffix."
        return model_file_name

    def _onnx_data_transformer(
        self,
        X: Union[pd.DataFrame, pd.Series],
        impute_values: Dict = None,
        force_overwrite: bool = False,
    ):
        """Apply onnx data transformer to data."""
        if self.framework in [Framework.TENSORFLOW, Framework.PYTORCH] or X is None:
            return X
        try:
            if hasattr(self, "onnx_data_preprocessor") and isinstance(
                self.onnx_data_preprocessor, ONNXTransformer
            ):
                X = self.onnx_data_preprocessor.transform(X=X)

            self.onnx_data_preprocessor = ONNXTransformer()
            X = self.onnx_data_preprocessor.fit_transform(
                X=X, impute_values=impute_values
            )
            if (
                os.path.exists(os.path.join(self.artifact_dir, ONNX_DATA_TRANSFORMER))
                and not force_overwrite
            ):
                raise ValueError(
                    f"{ONNX_DATA_TRANSFORMER} already exists. "
                    "Set `force_overwrite` to True if you wish to overwrite."
                )
            else:
                try:
                    self.onnx_data_preprocessor.save(
                        os.path.join(self.artifact_dir, ONNX_DATA_TRANSFORMER)
                    )
                except Exception as e:
                    logger.error(
                        f"Unable to serialize the data transformer due to: {e}."
                    )
                    raise e
        except Exception as e:
            logger.warn(f"Onnx Data Transformation was unsuccessful with error: {e}")
            raise e
        return X

    def prepare(
        self,
        inference_conda_env: str = None,
        inference_python_version: str = None,
        training_conda_env: str = None,
        training_python_version: str = None,
        model_file_name: str = None,
        as_onnx: bool = False,
        initial_types: List[Tuple] = None,
        force_overwrite: bool = False,
        namespace: str = CONDA_BUCKET_NS,
        use_case_type: str = None,
        X_sample: Union[list, tuple, pd.DataFrame, pd.Series, np.ndarray] = None,
        y_sample: Union[list, tuple, pd.DataFrame, pd.Series, np.ndarray] = None,
        training_script_path: str = None,
        training_id: str = _TRAINING_RESOURCE_ID,
        ignore_pending_changes: bool = True,
        max_col_num: int = DATA_SCHEMA_MAX_COL_NUM,
        **kwargs: Dict,
    ) -> None:
        """Prepare and save the score.py, serialized model and runtime.yaml file.

        Parameters
        ----------
        inference_conda_env: (str, optional). Defaults to None.
            Can be either slug or object storage path of the conda pack.
            You can only pass in slugs if the conda pack is a service pack.
        inference_python_version: (str, optional). Defaults to None.
            Python version which will be used in deployment.
        training_conda_env: (str, optional). Defaults to None.
            Can be either slug or object storage path of the conda pack.
            You can only pass in slugs if the conda pack is a service pack.
            If `training_conda_env` is not provided, `training_conda_env` will
            use the same value of `training_conda_env`.
        training_python_version: (str, optional). Defaults to None.
            Python version used during training.
        model_file_name: (str).
            Name of the serialized model.
        as_onnx: (bool, optional). Defaults to False.
            Whether to serialize as onnx model.
        initial_types: (list[Tuple], optional).
            Defaults to None. Only used for SklearnModel, LightGBMModel and XGBoostModel.
            Each element is a tuple of a variable name and a type.
            Check this link http://onnx.ai/sklearn-onnx/api_summary.html#id2 for
            more explanation and examples for `initial_types`.
        force_overwrite: (bool, optional). Defaults to False.
            Whether to overwrite existing files.
        namespace: (str, optional).
            Namespace of region. This is used for identifying which region the service pack
            is from when you pass a slug to inference_conda_env and training_conda_env.
        use_case_type: str
            The use case type of the model. Use it through UserCaseType class or string provided in `UseCaseType`. For
            example, use_case_type=UseCaseType.BINARY_CLASSIFICATION or use_case_type="binary_classification". Check
            with UseCaseType class to see all supported types.
        X_sample: Union[list, tuple, pd.Series, np.ndarray, pd.DataFrame]. Defaults to None.
            A sample of input data that will be used to generate input schema.
        y_sample: Union[list, tuple, pd.Series, np.ndarray, pd.DataFrame]. Defaults to None.
            A sample of output data that will be used to generate output schema.
        training_script_path: str. Defaults to None.
            Training script path.
        training_id: (str, optional). Defaults to value from environment variables.
            The training OCID for model. Can be notebook session or job OCID.
        ignore_pending_changes: bool. Defaults to False.
            whether to ignore the pending changes in the git.
        max_col_num: (int, optional). Defaults to utils.DATA_SCHEMA_MAX_COL_NUM.
            Do not generate the input schema if the input has more than this
            number of features(columns).
        kwargs:
            impute_values: (dict, optional).
                The dictionary where the key is the column index(or names is accepted
                for pandas dataframe) and the value is the impute value for the corresponding column.

        Raises
        ------
        FileExistsError: when files already exist but `force_overwrite` is False.
        ValueError: when `inference_python_version` is not provided, but also cannot be found through manifest file.

        Returns
        -------
        None
            Nothing.
        """
        # Populate properties from args and kwargs.
        # empty values will be ignored.
        locals_dict = _extract_locals(locals())
        locals_dict.pop("training_id", None)
        self.properties.with_dict(locals_dict)

        if training_id != _TRAINING_RESOURCE_ID:
            self.properties.training_id = training_id
        elif not self.properties.training_id:
            self.properties.training_id = _TRAINING_RESOURCE_ID

        if not self.properties.inference_conda_env:
            raise ValueError("`inference_conda_env` must be specified.")

        self._as_onnx = as_onnx
        self.model_file_name = self._handle_model_file_name(
            as_onnx=as_onnx, model_file_name=model_file_name
        )
        os.makedirs(self.artifact_dir, exist_ok=True)
        self.model_artifact = ModelArtifact(
            artifact_dir=self.artifact_dir,
            model_file_name=self.model_file_name,
        )

        self.runtime_info = self.model_artifact.prepare_runtime_yaml(
            inference_conda_env=self.properties.inference_conda_env,
            inference_python_version=self.properties.inference_python_version,
            training_conda_env=self.properties.training_conda_env,
            training_python_version=self.properties.training_python_version,
            force_overwrite=force_overwrite,
            namespace=namespace,
            bucketname=DEFAULT_CONDA_BUCKET_NAME,
        )

        self._summary_status.update_status(
            detail="Generated runtime.yaml", status=ModelState.DONE.value
        )

        if as_onnx:
            X_sample = self._onnx_data_transformer(
                X_sample,
                impute_values=kwargs.pop("impute_values", {}),
                force_overwrite=force_overwrite,
            )
        try:
            self.serialize_model(
                as_onnx=as_onnx,
                force_overwrite=force_overwrite,
                initial_types=initial_types,
                X_sample=X_sample,
                **kwargs,
            )
            self._summary_status.update_status(
                detail="Serialized model", status=ModelState.DONE.value
            )
        except SerializeModelNotImplementedError as e:
            if not os.path.exists(
                os.path.join(self.artifact_dir, self.model_file_name)
            ):
                self._summary_status.update_action(
                    detail="Serialized model",
                    action=f"Model is not automatically serialized. Serialize the model as `{self.model_file_name}` and save to the {self.artifact_dir}.",
                )
                self._summary_status.update_status(
                    detail="Serialized model", status=ModelState.NEEDSACTION.value
                )
                logger.warning(
                    f"{self.model_file_name} not found in {self.artifact_dir}. "
                    f"Save the serialized model under {self.artifact_dir}."
                )
                self._summary_status.update_action(
                    detail="Generated score.py",
                    action="`load_model` is not automatically generated. Finish implementing it and call .verify to check if it works.",
                )
        except Exception as e:
            raise e

        self._summary_status.update_status(
            detail="Generated runtime.yaml", status=ModelState.DONE.value
        )
        if as_onnx:
            jinja_template_filename = "score_onnx_new"
        else:
            if self.framework:
                jinja_template_filename = "score_" + self.framework
            else:
                jinja_template_filename = (
                    "score-pkl" if self._serialize else "score_generic"
                )

        self.model_artifact.prepare_score_py(jinja_template_filename)
        self._summary_status.update_status(
            detail="Generated score.py", status=ModelState.DONE.value
        )

        self.populate_metadata(
            use_case_type=use_case_type,
            X_sample=X_sample,
            y_sample=y_sample,
            training_script_path=self.properties.training_script_path,
            training_id=self.properties.training_id,
            ignore_pending_changes=ignore_pending_changes,
            max_col_num=max_col_num,
        )
        self._summary_status.update_status(
            detail="Populated metadata(Custom, Taxonomy and Provenance)",
            status=ModelState.DONE.value,
        )
        self._summary_status.update_status(
            detail="Local tested .predict from score.py",
            status=ModelState.AVAILABLE.value,
        )
        self._summary_status.update_status(
            detail="Conducted Introspect Test", status=ModelState.AVAILABLE.value
        )
        self._summary_status.update_status(
            detail="Uploaded artifact to model catalog",
            status=ModelState.AVAILABLE.value,
        )

    def verify(self, data: Any) -> Dict[str, Any]:
        """test if deployment works in local environment.

        Parameters
        ----------
        data: Any.
            Data used to test if deployment works in local environment.

        Returns
        -------
        Dict
            A dictionary which contains prediction results.
        """
        url = f"http://127.0.0.1:8000/predict"
        headers = {
            "Content-Type": "application/json",
            "Accept": "*/*",
        }
        data_dict = self._handle_input_data(data)
        req = requests.Request("POST", url, json=data_dict, headers=headers)
        prepared_req = req.prepare()
        self.model_artifact.reload()
        prediction = self.model_artifact.predict(json.loads(prepared_req.body))

        try:
            req = requests.Request("POST", url, json=prediction, headers=headers)
        except:
            raise TypeError(
                "The prediction result is not json serializable. "
                "Please modify the score.py."
            )

        self._summary_status.update_status(
            detail="Local tested .predict from score.py", status=ModelState.DONE.value
        )
        return prediction

    def _serialize_input(
        self,
        data: Union[
            Dict,
            str,
            List,
            np.ndarray,
            pd.core.series.Series,
            pd.core.frame.DataFrame,
        ],
    ) -> Union[Dict, List, str]:
        """Returns serializable input data.

        Parameters
        ----------
        data: Union[Dict, str, list, numpy.ndarray,
        pd.core.series.Series,  pd.core.frame.DataFrame]
        data_type: (Any, otional). Defaults to None.
            Type of the data. If not provided, it will be checked against data.

        Returns
        -------
        Union[Dict, List, str]
            Serializable input data.

        Raises
        ------
        TypeError
            If provided data type is not supported.
        """
        raise SerializeInputNotImplementedError(
            "`_serialize_input` is not implemented."
        )

    def introspect(self) -> pd.DataFrame:
        """Conducts instrospection.

        Returns
        -------
        pandas.DataFrame
            A pandas DataFrame which contains the instrospection results.
        """
        df = self._introspect()
        return df

    @classmethod
    def from_model_artifact(
        cls,
        uri: str,
        model_file_name: str,
        artifact_dir: str,
        auth: Optional[Dict] = None,
        force_overwrite: Optional[bool] = False,
        properties: Optional[ModelProperties] = None,
        **kwargs: dict,
    ) -> "GenericModel":
        """Loads model from a folder, or zip/tar archive.

        Parameters
        ----------
        uri: str
            The folder path, ZIP file path, or TAR file path. It could contain a
            seriliazed model(required) as well as any files needed for deployment including:
            serialized model, runtime.yaml, score.py and etc. The content of the folder will be
            copied to the `artifact_dir` folder.
        model_file_name: str
            The serialized model file name.
        artifact_dir: str
            The artifact directory to store the files needed for deployment.
        auth: (Dict, optional). Defaults to None.
            The default authetication is set using `ads.set_auth` API. If you need to override the
            default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
            authentication signer and kwargs required to instantiate IdentityClient object.
        force_overwrite: (bool, optional). Defaults to False.
            Whether to overwrite existing files or not.
        properties: (ModelProperties, optional). Defaults to None.
            ModelProperties object required to save and deploy model.

        Returns
        -------
        GenericModel
            An instance of `GenericModel` class.

        Raises
        ------
        ValueError
            If `model_file_name` not provided.
        """
        if not model_file_name or not isinstance(model_file_name, str):
            raise ValueError("The `model_file_name` must be provided.")

        local_vars = _extract_locals(locals())
        properties = properties or ModelProperties()
        properties.with_dict(local_vars)

        auth = auth or authutil.default_signer()
        model_artifact = ModelArtifact.from_uri(
            uri=uri,
            artifact_dir=artifact_dir,
            model_file_name=model_file_name,
            force_overwrite=force_overwrite,
            auth=auth,
        )

        model = cls(
            estimator=model_artifact.model,
            artifact_dir=artifact_dir,
            auth=auth,
            properties=properties,
        )
        model.model_file_name = model_file_name
        model.model_artifact = model_artifact
        model.reload()
        model._summary_status.update_status(
            detail="Generated score.py", status=ModelState.DONE.value
        )
        model._summary_status.update_status(
            detail="Generated runtime.yaml", status=ModelState.DONE.value
        )
        model._summary_status.update_status(
            detail="Serialized model", status=ModelState.DONE.value
        )
        model._summary_status.update_action(
            detail="Populated metadata(Custom, Taxonomy and Provenance)",
            action=f"Call .populate_metadata() to populate metadata.",
        )
        return model

    @classmethod
    def from_model_catalog(
        cls,
        model_id: str,
        model_file_name: str,
        artifact_dir: str,
        auth: Optional[Dict] = None,
        force_overwrite: Optional[bool] = False,
        properties: Optional[Union[ModelProperties, Dict]] = None,
        **kwargs,
    ) -> "GenericModel":
        """Loads model from model catalog.

        Parameters
        ----------
        model_id: str
            The model OCID.
        model_file_name: (str)
            The name of the serialized model.
        artifact_dir: str
            The artifact directory to store the files needed for deployment.
            Will be created if not exists.
        auth: (Dict, optional). Defaults to None.
            The default authetication is set using `ads.set_auth` API. If you need to override the
            default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
            authentication signer and kwargs required to instantiate IdentityClient object.
        force_overwrite: (bool, optional). Defaults to False.
            Whether to overwrite existing files or not.
        properties: (ModelProperties, optional). Defaults to None.
            ModelProperties object required to save and deploy model.
        kwargs:
            compartment_id : (str, optional)
                Compartment OCID. If not specified, the value will be taken from the environment variables.
            timeout : (int, optional). Defaults to 10 seconds.
                The connection timeout in seconds for the client.

        Returns
        -------
        GenericModel
            An instance of GenericModel class.
        """
        local_vars = _extract_locals(locals())
        properties = properties or ModelProperties()
        properties.with_dict(local_vars)
        properties.compartment_id = properties.compartment_id or _COMPARTMENT_OCID

        auth = auth or authutil.default_signer()
        artifact_dir = os.path.abspath(os.path.expanduser(artifact_dir))

        model_catalog = ModelCatalog(
            compartment_id=properties.compartment_id,
            ds_client_auth=auth,
            identity_client_auth=auth,
            timeout=kwargs.pop("timeout", None),
        )

        model_catalog._download_artifacts(model_id, artifact_dir, force_overwrite)
        oci_model = model_catalog.get_model(model_id)

        result_model = cls.from_model_artifact(
            uri=artifact_dir,
            model_file_name=model_file_name,
            artifact_dir=artifact_dir,
            auth=auth,
            force_overwrite=force_overwrite,
            properties=properties,
        )

        # Populate metadata
        result_model.metadata_custom = oci_model.metadata_custom
        result_model.metadata_taxonomy = oci_model.metadata_taxonomy
        result_model.schema_input = oci_model.schema_input
        result_model.schema_output = oci_model.schema_output
        result_model.metadata_provenance = oci_model.provenance_metadata
        result_model.model_catalog = model_catalog
        result_model._summary_status.update_status(
            detail="Populated metadata(Custom, Taxonomy and Provenance)",
            status=ModelState.DONE.value,
        )
        result_model._summary_status.update_action(
            detail="Populated metadata(Custom, Taxonomy and Provenance)", action=""
        )
        return result_model

    def reload(self) -> None:
        """Reloads the model artifact files: `score.py` and the `runtime.yaml`.

        Returns
        -------
        None
            Nothing.
        """
        # reload the score.py
        self.model_artifact.reload()

        # reload runtime.yaml
        runtime_yaml_file = os.path.join(self.artifact_dir, "runtime.yaml")
        if not os.path.exists(runtime_yaml_file):
            raise FileNotFoundError(
                f"`runtime.yaml` does not exist in {self.artifact_dir}. "
                "Use `RuntimeInfo` class to populate it."
            )
        self.runtime_info = RuntimeInfo.from_yaml(uri=runtime_yaml_file)

    def save(
        self,
        display_name: Optional[str] = None,
        description: Optional[str] = None,
        freeform_tags: Optional[dict] = None,
        defined_tags: Optional[dict] = None,
        ignore_introspection: Optional[bool] = False,
        **kwargs,
    ) -> None:
        """Saves model artifacts to the model catalog.

        Parameters
        ----------
        display_name: (str, optional). Defaults to None.
            The name of the model.
        description: (str, optional). Defaults to None.
            The description of the model.
        freeform_tags : Dict(str, str), Defaults to None.
            Freeform tags for the model.
        defined_tags : (Dict(str, dict(str, object)), optional). Defaults to None.
            Defined tags for the model.
        ignore_introspection: (bool, optional). Defaults to None.
            Determine whether to ignore the result of model introspection or not.
            If set to True, the save will ignore all model introspection errors.
        kwargs:
            project_id: (str, optional).
                Project OCID. If not specified, the value will be taken either
                from the environment variables or model properties.
            compartment_id : (str, optional).
                Compartment OCID. If not specified, the value will be taken either
                from the environment variables or model properties.
            timeout: (int, optional). Defaults to 10 seconds.
                The connection timeout in seconds for the client.
        Raises
        ------
        RuntimeInfoInconsistencyError
            When `.runtime_info` is not synched with runtime.yaml file.

        Returns
        -------
        None
            Nothing
        """
        # populates properties from args and kwargs. Empty values will be ignored.
        self.properties.with_dict(_extract_locals(locals()))
        self.properties.compartment_id = (
            self.properties.compartment_id or _COMPARTMENT_OCID
        )
        self.properties.project_id = self.properties.project_id or PROJECT_OCID

        # check if the runtime_info sync with the runtime.yaml.
        runtime_file_path = os.path.join(self.artifact_dir, "runtime.yaml")
        runtime_info_from_yaml = RuntimeInfo.from_yaml(uri=runtime_file_path)
        if self.runtime_info != runtime_info_from_yaml:
            raise RuntimeInfoInconsistencyError(
                "`.runtime_info` does not sync with runtime.yaml file. Call "
                "`.runtime_info.save()` if you updated `runtime_info`. "
                "Call `.reload()` if you updated runtime.yaml file."
            )
        # reload to check if load_model works in score.py, i.e.
        # whether the model file has been serialized, and whether it can be loaded
        # successfully.
        self.reload()
        if not ignore_introspection:
            self._introspect()
            if self._introspect.status == TEST_STATUS.NOT_PASSED:
                msg = (
                    "Model introspection not passed. "
                    "See the table below for "
                    "more detailed information and follow the messages to fix it, "
                    "otherwise you might come across errors during model deployment. "
                    "To save model artifacts ignoring "
                    "introspection use .save(ignore_introspection=True...)."
                )
                logger.error(msg)
                display(self._introspect.to_dataframe())

                self._summary_status.update_status(
                    detail="Conducted Introspect Test", status="Failed"
                )
                self._summary_status.update_action(
                    detail="Conducted Introspect Test",
                    action=f"Check introspect table for the reason why it failed.",
                )
                raise IntrospectionNotPassed()
            else:
                self._summary_status.update_status(
                    detail="Conducted Introspect Test", status=ModelState.DONE.value
                )
                self._summary_status.update_action(
                    detail="Conducted Introspect Test", action=""
                )

        # upload model artifact to model catalog.
        self.model_catalog = ModelCatalog(
            compartment_id=self.properties.compartment_id,
            ds_client_auth=self.auth,
            identity_client_auth=self.auth,
            timeout=kwargs.pop("timeout", None),
        )
        oci_model = self.model_catalog.upload_model(
            model_artifact=self,
            provenance_metadata=self.metadata_provenance._to_oci_metadata(),
            project_id=self.properties.project_id,
            display_name=display_name,
            description=description,
            freeform_tags=freeform_tags,
            defined_tags=defined_tags,
        )
        self.model_id = oci_model.id

        self._summary_status.update_status(
            detail="Uploaded artifact to model catalog", status=ModelState.DONE.value
        )
        self._summary_status.update_status(
            detail="Deployed the model", status=ModelState.AVAILABLE.value
        )

        return self.model_id

    def _get_files(self):
        """List out all the file names under the artifact_dir.

        Returns
        -------
        List
            List of the files in the artifact_dir.
        """
        return get_files(self.artifact_dir)

    def deploy(
        self,
        wait_for_completion: Optional[bool] = True,
        display_name: Optional[str] = None,
        description: Optional[str] = None,
        deployment_instance_shape: Optional[str] = None,
        deployment_instance_count: Optional[int] = None,
        deployment_bandwidth_mbps: Optional[int] = None,
        deployment_log_group_id: Optional[str] = None,
        deployment_access_log_id: Optional[str] = None,
        deployment_predict_log_id: Optional[str] = None,
        **kwargs: Dict,
    ) -> ModelDeployment:
        """
        Deploys a model. The model needs to be saved to the model catalog at first.

        Parameters
        ----------
        wait_for_completion : (bool, optional). Defaults to True.
            Flag set for whether to wait for deployment to complete before proceeding.
        display_name: (str, optional). Defaults to None.
            The name of the model.
        description: (str, optional). Defaults to None.
            The description of the model.
        deployment_instance_shape: (str, optional). Default to `VM.Standard2.1`.
            The shape of the instance used for deployment.
        deployment_instance_count: (int, optional). Defaults to 1.
            The number of instance used for deployment.
        deployment_bandwidth_mbps: (int, optional). Defaults to 10.
            The bandwidth limit on the load balancer in Mbps.
        deployment_log_group_id: (str, optional). Defaults to None.
            The oci logging group id. The access log and predict log share the same log group.
        deployment_access_log_id: (str, optional). Defaults to None.
            The access log OCID for the access logs. https://docs.oracle.com/en-us/iaas/data-science/using/model_dep_using_logging.htm
        deployment_predict_log_id: (str, optional). Defaults to None.
            The predict log OCID for the predict logs. https://docs.oracle.com/en-us/iaas/data-science/using/model_dep_using_logging.htm
        kwargs:
            project_id: (str, optional).
                Project OCID. If not specified, the value will be taken from the environment variables.
            compartment_id : (str, optional).
                Compartment OCID. If not specified, the value will be taken from the environment variables.
            max_wait_time : (int, optional). Defaults to 1200 seconds.
                Maximum amount of time to wait in seconds.
                Negative implies infinite wait time.
            poll_interval : (int, optional). Defaults to 60 seconds.
                Poll interval in seconds.

        Returns
        -------
        ModelDeployment
            The ModelDeployment instance.

        Raises
        ------
        ValueError
            If `model_id` is not specified.
        """
        # populates properties from args and kwargs. Empty values will be ignored.
        self.properties.with_dict(_extract_locals(locals()))
        # clears out project_id and compartment_id from kwargs, to prevent passing
        # these params to the deployment via kwargs.
        kwargs.pop("project_id", None)
        kwargs.pop("compartment_id", None)

        self.properties.compartment_id = (
            self.properties.compartment_id or _COMPARTMENT_OCID
        )
        self.properties.project_id = self.properties.project_id or PROJECT_OCID
        self.properties.deployment_instance_shape = (
            self.properties.deployment_instance_shape or MODEL_DEPLOYMENT_INSTANCE_SHAPE
        )
        self.properties.deployment_instance_count = (
            self.properties.deployment_instance_count or MODEL_DEPLOYMENT_INSTANCE_COUNT
        )
        self.properties.deployment_bandwidth_mbps = (
            self.properties.deployment_bandwidth_mbps or MODEL_DEPLOYMENT_BANDWIDTH_MBPS
        )

        self.model_deployment = None

        if not self.model_id:
            raise ValueError(
                "The model needs to be saved to the Model Catalog "
                "before it can be depployed."
            )

        if (
            self.properties.deployment_access_log_id
            or self.properties.deployment_predict_log_id
        ) and not self.properties.deployment_log_group_id:
            raise ValueError(
                "`deployment_log_group_id` needs to be specified. "
                "`deployment_access_log_id` and `deployment_predict_log_id` "
                "cannot be used without `deployment_log_group_id`."
            )

        model_deployment_properties = (
            ModelDeploymentProperties(self.model_id, config=self.auth)
            .with_prop("display_name", display_name)
            .with_prop("description", description)
            .with_prop("project_id", self.properties.project_id)
            .with_prop("compartment_id", self.properties.compartment_id)
            .with_instance_configuration(
                {
                    "INSTANCE_SHAPE": self.properties.deployment_instance_shape,
                    "INSTANCE_COUNT": self.properties.deployment_instance_count,
                    "bandwidth_mbps": self.properties.deployment_bandwidth_mbps,
                }
            )
        )

        # specifies the access log id
        if self.properties.deployment_access_log_id:
            model_deployment_properties.with_access_log(
                self.properties.deployment_log_group_id,
                self.properties.deployment_access_log_id,
            )

        # specifies the predict log id
        if self.properties.deployment_predict_log_id:
            model_deployment_properties.with_predict_log(
                self.properties.deployment_log_group_id,
                self.properties.deployment_predict_log_id,
            )

        self.model_deployment = ModelDeployer(config=self.auth).deploy(
            properties=model_deployment_properties,
            wait_for_completion=wait_for_completion,
            max_wait_time=kwargs.pop("max_wait_time", DEFAULT_WAIT_TIME),
            poll_interval=kwargs.pop("poll_interval", DEFAULT_POLL_INTERVAL),
            **kwargs,
        )
        self._summary_status.update_status(
            detail="Deployed the model",
            status=self.model_deployment.state.name.upper(),
        )
        return self.model_deployment

    def predict(self, data: Any) -> Dict[str, Any]:
        """Returns prediction of input data run against the model deployment endpoint.

        Parameters
        ----------
        data: Any
            JSON serializable data for the prediction for onnx models, for local serialization
            method, data can be the data types that each framework support.

        Returns
        -------
        Dict[str, Any]
            Dictionary with the predicted values.

        Raises
        ------
        NotActiveDeploymentError
            If model deployment process was not started or not finished yet.
        ValueError
            If `data` is empty or not JSON serializable.
        """
        if not self.model_deployment:
            raise ValueError("Use `deploy()` method to start model deployment.")

        current_state = self.model_deployment.state.name.upper()
        if current_state != ModelDeploymentState.ACTIVE.name:
            raise NotActiveDeploymentError(current_state)

        data_dict = self._handle_input_data(data)
        prediction = self.model_deployment.predict(data_dict)

        self._summary_status.update_status(
            detail="Called deployment predict endpoint", status=ModelState.DONE.value
        )
        return prediction

    def _handle_input_data(self, data: any):
        """For generic model where _serialize_input is not implemented, check
        if the data is json serializable. For other frameworks, serialize the
        data using _serialize_input and double check if the resulted data is
        json serializable.

        Parameters
        ----------
        data: (Any)
            data to be passed to model for prediction.

        Returns
        -------
        data
            Serialized data.
        """
        try:
            data_dict = self._serialize_input(data)
        except SerializeInputNotImplementedError as e:
            data_dict = data
            if not GenericModel._is_json_serializable(data):
                raise ValueError(
                    "`data` must be json serializable. "
                    "Serialize the provided input data first."
                )
        except Exception as e:
            raise e
        if not GenericModel._is_json_serializable(data_dict):
            raise ValueError(
                f"An error occurred in attempt to serialize the `data` to JSON. "
                "Check if the type of the data is supported."
            )
        return data_dict

    @staticmethod
    def _is_json_serializable(data: Any) -> bool:
        """Check is data input is json serialization.

        Parameters
        ----------
        data: (Any)
            data to be passed to model for prediction.

        Returns
        -------
        bool
            Whether data is json serializable.
        """
        result = True
        try:
            json.dumps(data)
        except:
            result = False
        return result

    def summary_status(self) -> pd.DataFrame:
        """A summary table of the current status.

        Returns
        -------
        pd.DataFrame
            The summary stable of the current status.
        """
        if self.model_file_name and not os.path.exists(
            os.path.join(self.artifact_dir, self.model_file_name)
        ):
            self._summary_status.update_action(
                detail="Serialized model",
                action=f"Model is not automatically serialized. Serialize the model as `{self.model_file_name}` and save to the {self.artifact_dir}.",
            )
            self._summary_status.update_status(
                detail="Serialized model", status=ModelState.NEEDSACTION.value
            )
        else:
            self._summary_status.update_action(detail="Serialized model", action="")
            if (
                ModelState.NEEDSACTION.value
                in self._summary_status.df.loc[
                    self._summary_status.df["Details"] == "Serialized model", "Status"
                ].values
            ):
                self._summary_status.update_status(
                    detail="Serialized model", status=ModelState.DONE.value
                )
        if (
            self._summary_status.df.loc[
                self._summary_status.df["Details"] == "Generated score.py",
                "Actions Needed",
            ].values
            != ""
        ):
            try:
                self.model_artifact.reload()
                self._summary_status.update_action(
                    detail="Generated score.py", action=""
                )
            except:
                pass

        if self.model_deployment:
            self._summary_status.update_status(
                detail="Deployed the model",
                status=self.model_deployment.state.name.upper(),
            )

            if self.model_deployment.state == ModelDeploymentState.ACTIVE:
                self._summary_status.update_status(
                    detail="Called deployment predict endpoint",
                    status=ModelState.AVAILABLE.value,
                )
            elif (
                self.model_deployment.state
                and self.model_deployment.state != ModelDeploymentState.ACTIVE
            ):
                self._summary_status.update_status(
                    detail="Called deployment predict endpoint",
                    status=ModelState.NOTAVAILABLE.value,
                )

        return self._summary_status.df.set_index(["Step", "Status", "Details"])

    def delete_deployment(self, wait_for_completion: bool = False):
        """Deletes the current deployment.

        Parameters
        ----------
        wait_for_completion: (bool, optional). Defaults to False.
            Whether to wait till completion.

        Raises
        ------
        ValueError: if there is not deployment attached yet.
        """
        if not self.model_deployment:
            raise ValueError("Use `deploy()` method to start model deployment.")
        self.model_deployment.delete(wait_for_completion=wait_for_completion)

    def _delete_model(self, model_id: str = None):
        """Deletes the remote model artifact on the model catalog.

        Parameters
        ----------
        model_id: (str, optional). Defaults to current saved model artfact id.
            model artifact id to be deleted.

        Raises
        ------
        ValueError: if model id cannot be found.
        """
        if model_id:
            model_id = model_id
        elif hasattr(self, "model_id"):
            model_id = self.model_id
        else:
            raise ValueError(
                "Cannot find model to delete. `model_id` needs to be provided."
            )
        self.model_catalog.delete_model(model=model_id)


class ModelState(Enum):
    DONE = "Done"
    AVAILABLE = "Available"
    NOTAVAILABLE = "Not Available"
    NEEDSACTION = "Needs Action"


class SummaryStatus:
    """SummaryStatus class which track the status of the Model frameworks."""

    def __init__(self):
        summary_data = [
            ["initiate", "Initiated the model", ModelState.DONE.value, ""],
            ["prepare()", "Generated runtime.yaml", ModelState.AVAILABLE.value, ""],
            ["prepare()", "Generated score.py", ModelState.AVAILABLE.value, ""],
            ["prepare()", "Serialized model", ModelState.AVAILABLE.value, ""],
            [
                "prepare()",
                "Populated metadata(Custom, Taxonomy and Provenance)",
                ModelState.AVAILABLE.value,
                "",
            ],
            [
                "verify()",
                "Local tested .predict from score.py",
                ModelState.NOTAVAILABLE.value,
                "",
            ],
            ["save()", "Conducted Introspect Test", ModelState.NOTAVAILABLE.value, ""],
            [
                "save()",
                "Uploaded artifact to model catalog",
                ModelState.NOTAVAILABLE.value,
                "",
            ],
            ["deploy()", "Deployed the model", ModelState.NOTAVAILABLE.value, ""],
            [
                "predict()",
                "Called deployment predict endpoint",
                ModelState.NOTAVAILABLE.value,
                "",
            ],
        ]
        self.df = pd.DataFrame(
            summary_data, columns=["Step", "Details", "Status", "Actions Needed"]
        )

    def update_status(self, detail: str, status: str) -> None:
        """Updates the status of the summary status table of the corresponding detail.

        Parameters
        ----------
        detail: (str)
            value of the detail in the Details column. Used to locate which row to update.
        status: (str)
            new status to be updated for the row specified by detail.

        Returns
        -------
        None
            Nothing.
        """
        self.df.loc[
            self.df["Details"] == detail,
            "Status",
        ] = status

    def update_action(self, detail: str, action: str) -> None:
        """Updates the action of the summary status table of the corresponding detail.

        Parameters
        ----------
        detail: (str)
            Value of the detail in the Details column. Used to locate which row to update.
        status: (str)
            New status to be updated for the row specified by detail.

        Returns
        -------
        None
            Nothing.
        """
        self.df.loc[
            self.df["Details"] == detail,
            "Actions Needed",
        ] = action
