import json
import numpy as np
import logging
import onnxruntime as rt
import os
import sys
import pandas as pd
from functools import lru_cache
from scipy.sparse.csr import csr_matrix
from typing import Union, Dict


model_name = '{{model_file_name}}'
transformer_name = 'onnx_data_transformer.json'


"""
   Inference script. This script is used for prediction by scoring server when schema is known.
"""


@lru_cache(maxsize=10)
def load_model(model_file_name=model_name):
    """
    Loads model from the serialized format

    Returns
    -------
    model:  an onnxruntime session instance
    """
    model_dir = os.path.dirname(os.path.realpath(__file__))
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    contents = os.listdir(model_dir)
    if model_file_name in contents:
        print(f'Start loading {model_file_name} from model directory {model_dir} ...')
        loaded_model = rt.InferenceSession(os.path.join(model_dir, model_file_name))

        print("Model is successfully loaded.")
        return loaded_model
    else:
        raise Exception(f'{model_file_name} is not found in model directory {model_dir}')


def predict(data, model=load_model()):
    """
    Returns prediction given the model and data to predict

    Parameters
    ----------
    model: Model session instance returned by load_model API
    data: Data format as expected by the onnxruntime API

    Returns
    -------
    predictions: Output from scoring server
        Format: {'prediction':output from model.predict method}

    """
    from pandas import read_json, DataFrame
    from io import StringIO
    X = read_json(StringIO(data)) if isinstance(data, str) else DataFrame.from_dict(data)
    model_dir = os.path.dirname(os.path.realpath(__file__))
    contents = os.listdir(model_dir)
    # Note: User may need to edit this
    if transformer_name in contents:
        onnx_data_transformer = ONNXTransformer.load(os.path.join(model_dir, transformer_name))
        X = onnx_data_transformer.transform(X)
    else:
        onnx_data_transformer = None
    onnx_transformed_rows = []
    for name, row in X.iterrows():
        onnx_transformed_rows.append(list(row))
    input_data = {model.get_inputs()[0].name: onnx_transformed_rows}
    pred = model.run(None, input_data)
    return {'prediction':pred[0].tolist()}


class ONNXTransformer(object):
    """
    This is a transformer to convert X [pandas.Dataframe, equivalent] data into Onnx
    readable dtypes and formats. It is Serializable, so it can be reloaded at another time.


    Examples
    --------
    >>> from ads.model.transformer.onnx_transformer import ONNXTransformer
    >>> onnx_data_transformer = ONNXTransformer()
    >>> train_transformed = onnx_data_transformer.fit_transform(train.X, {"column_name1": "impute_value1", "column_name2": "impute_value2"}})
    >>> test_transformed = onnx_data_transformer.transform(test.X)
    """

    def __init__(self):
        self.impute_values = {}
        self.dtypes = None
        self._fitted = False

    @staticmethod
    def _handle_dtypes(X: Union[pd.DataFrame, pd.Series]):
        """handle the dtypes for pandas dataframe and pandas Series.

        Parameters
        ----------
        X : Union[pd.DataFrame, pd.Series]
            The Dataframe for the training data

        Returns
        -------
        Union[pd.DataFrame, pd.Series]
            The transformed(numerical values are cast to float32) X data
        """
        # Data type cast could be expensive doing it in a for loop
        # Especially with wide datasets
        # So cast the numerical columns first, without loop
        # Then impute missing values
        if isinstance(X, pd.Series):
            series_name = X.name if X.name else 0
            _X = X.to_frame()
            _X = ONNXTransformer._handle_dtypes_dataframe(_X)[series_name]
        elif isinstance(X, pd.DataFrame):
            _X = ONNXTransformer._handle_dtypes_dataframe(X)
        else:
            raise NotImplemented(f"{type(X)} is not supported.")
        return _X

    @staticmethod
    def _handle_dtypes_dataframe(X: pd.DataFrame):
        """handle the dtypes for pandas dataframe.

        Parameters
        ----------
        X : pandas.DataFrame
            The Dataframe for the training data

        Returns
        -------
        pandas.DataFrame
            The transformed X data
        """
        dict_astype = {}
        for k, v in zip(X.columns, X.dtypes):
            if "int" in str(v) or "float" in str(v):
                dict_astype[k] = "float32"
        _X = X.astype(dict_astype)
        return _X

    def fit(self, X: Union[pd.DataFrame, pd.Series], impute_values: Dict = {}):
        """
        Fits the OnnxTransformer on the dataset
        Parameters
        ----------
        X : Union[pandas.DataFrame, pandas.Series]
            The Dataframe for the training data

        Returns
        -------
        Self: ads.Model
            The fitted estimator
        """
        _X = ONNXTransformer._handle_dtypes(X)
        self.dtypes = _X.dtypes
        self.impute_values = impute_values
        self._fitted = True
        return self

    def transform(self, X: Union[pd.DataFrame, pd.Series]):
        """
        Transforms the data for the OnnxTransformer.

        Parameters
        ----------
        X: Union[pandas.DataFrame, pandas.Series]
            The Dataframe for the training data

        Returns
        -------
        Union[pandas.DataFrame, pandas.Series]
            The transformed X data
        """
        assert self._fitted, "Call fit_transform first!"
        _X = X.astype(self.dtypes)
        logging.warning("Numerical values in `X` are cast to float32.")
        _X = ONNXTransformer._handle_missing_value(_X, impute_values=self.impute_values)
        return _X

    @staticmethod
    def _handle_missing_value(X: Union[pd.DataFrame, pd.Series], impute_values: Dict):
        """Impute missing values in X according to impute_values.

        Parameters
        ----------
        X: Union[pandas.DataFrame, pandas.Series]
            The Dataframe for the training data

        Raises
        ------
        Exception if X has only one dim, but imputed_values has multiple values.
        NotImplemented if X has the data type that is not supported.

        Returns
        -------
        Union[pandas.DataFrame, pd.Series]
            The transformed X data
        """
        if isinstance(X, pd.DataFrame):
            X = ONNXTransformer._handle_missing_value_dataframe(
                X, impute_values=impute_values
            )
        elif isinstance(X, pd.Series):
            X = X.replace(r"^\s*$", np.NaN, regex=True)
            if len(impute_values.keys()) == 1:
                for key, val in impute_values.items():
                    X = X.fillna(val)
            else:
                raise Exception(
                    "Multiple imputed values are provided, but `X` has only one dim."
                )
        else:
            raise NotImplemented(
                f"{type(X)} is not supported. Supported types includes pd.DataFrame, and pd.Series."
            )
        return X

    @staticmethod
    def _handle_missing_value_dataframe(X: pd.DataFrame, impute_values: Dict):
        for idx, val in impute_values.items():
            if isinstance(idx, int):
                X.iloc[:, idx] = (
                    X.iloc[:, idx].replace(r"^\s*$", np.NaN, regex=True).fillna(val)
                )
            else:
                X.loc[:, idx] = (
                    X.loc[:, idx].replace(r"^\s*$", np.NaN, regex=True).fillna(val)
                )
        return X

    def fit_transform(
        self, X: Union[pd.DataFrame, pd.Series], impute_values: Dict = {}
    ):
        """
        Fits, then transforms the data
        Parameters
        ----------
        X: Union[pandas.DataFrame, pandas.Series]
            The Dataframe for the training data

        Returns
        -------
        Union[pandas.DataFrame, pandas.Series]
            The transformed X data
        """
        return self.fit(X, impute_values).transform(X)

    def save(self, filename, **kwargs):
        """
        Saves the Onnx model to disk
        Parameters
        ----------
        filename: Str
            The filename location for where the model should be saved

        Returns
        -------
        filename: Str
            The filename where the model was saved
        """
        export_dict = {
            "impute_values": {
                "value": self.impute_values,
                "dtype": str(type(self.impute_values)),
            },
            "dtypes": {
                "value": {
                    "index": list(self.dtypes.index),
                    "values": [str(val) for val in self.dtypes.values],
                }
                if self.dtypes is not None
                else {},
                "dtype": str(type(self.dtypes)),
            },
            "_fitted": {"value": self._fitted, "dtype": str(type(self._fitted))},
        }

        with open(filename, "w") as f:
            json.dump(export_dict, f, sort_keys=True, indent=4, separators=(",", ": "))
        return filename

    @staticmethod
    def load(filename, **kwargs):
        """
        Loads the Onnx model to disk
        Parameters
        ----------
        filename: Str
            The filename location for where the model should be loaded

        Returns
        -------
        onnx_transformer: ONNXTransformer
            The loaded model
        """
        # Make sure you have  pandas, numpy, and sklearn imported
        with open(filename, "r") as f:
            export_dict = json.load(f)

        onnx_transformer = ONNXTransformer()

        for key in export_dict.keys():
            if key not in ["impute_values", "dtypes"]:
                try:
                    setattr(onnx_transformer, key, export_dict[key]["value"])
                except Exception as e:
                    print(
                        f"Warning: Failed to reload {key} from {filename} to OnnxTransformer."
                    )
                    raise e
        onnx_transformer.dtypes = pd.Series(
            data=[np.dtype(val) for val in export_dict["dtypes"]["value"]["values"]],
            index=export_dict["dtypes"]["value"]["index"],
        )
        onnx_transformer.impute_values = export_dict["impute_values"]["value"]
        return onnx_transformer
