# SPDX-License-Identifier: Apache-2.0

import pickle
import os
import warnings
import numpy
import packaging.version as pv
from onnx.defs import onnx_opset_version
from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER
from ..convert.common.data_types import FloatTensorType
from .utils_backend import compare_backend, extract_options, evaluate_condition, is_backend_enabled


TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version())


def dump_data_and_model(data, model, onnx=None, basename="model", folder=None,
                        inputs=None, backend="onnxruntime", context=None,
                        allow_failure=None, verbose=False):
    """
    Saves data with pickle, saves the model with pickle and *onnx*,
    runs and saves the predictions for the given model.
    This function is used to test a backend (runtime) for *onnx*.

    :param data: any kind of data
    :param model: any model
    :param onnx: *onnx* model or *None* to use *onnxmltools* to convert it
        only if the model accepts one float vector
    :param basemodel: three files are writen ``<basename>.data.pkl``,
        ``<basename>.model.pkl``, ``<basename>.model.onnx``
    :param folder: files are written in this folder,
        it is created if it does not exist, if *folder* is None,
        it looks first in environment variable ``ONNXTESTDUMP``,
        otherwise, it is placed into ``'tests'``.
    :param inputs: standard type or specific one if specified, only used is
        parameter *onnx* is None
    :param backend: backend used to compare expected output and runtime output.
        Two options are currently supported: None for no test,
        `'onnxruntime'` to use module *onnxruntime*.
    :param context: used if the model contains a custom operator
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :param verbose: prints more information when it fails
    :return: the created files

    Some convention for the name,
    *Bin* for a binary classifier, *Mcl* for a multiclass
    classifier, *Reg* for a regressor, *MRg* for a multi-regressor.
    The name can contain some flags. Expected outputs refer to the
    outputs computed with the original library, computed outputs
    refer to the outputs computed with a ONNX runtime.

    * ``-CannotLoad``: the model can be converted but the runtime cannot load it
    * ``-Dec3``: compares expected and computed outputs up to 3 decimals (5 by default)
    * ``-Dec4``: compares expected and computed outputs up to 4 decimals (5 by default)
    * ``-NoProb``: The original models computed probabilites for two classes *size=(N, 2)*
      but the runtime produces a vector of size *N*, the test will compare the second column
      to the column
    * ``-OneOff``: the ONNX runtime cannot computed the prediction for several inputs,
      it must be called for each of them
      and computed output.
    * ``-Out0``: only compares the first output on both sides
    * ``-Reshape``: merges all outputs into one single vector and resizes it before comparing
    * ``-SkipDim1``: before comparing expected and computed output,
      arrays with a shape like *(2, 1, 2)* becomes *(2, 2)*

    If the *backend* is not None, the function either raises an exception
    if the comparison between the expected outputs and the backend outputs
    fails or it saves the backend output and adds it to the results.
    """
    runtime_test = dict(model=model, data=data)

    if folder is None:
        folder = os.environ.get('ONNXTESTDUMP', 'tests/temp')
    if not os.path.exists(folder):
        os.makedirs(folder)

    if hasattr(model, "predict"):
        import lightgbm
        import xgboost
        if isinstance(model, lightgbm.Booster):
            # LightGBM Booster
            model_dict = model.dump_model()
            if model_dict['objective'].startswith('binary'):
                score = model.predict(data)
                if len(score.shape) < 2 or score.shape[1] == 1:
                    score = score.ravel()
                    score = numpy.vstack([1-score, score]).T
                prediction = [score[:, 1] > 0.5, score]
            elif model_dict['objective'].startswith('multiclass'):
                score = model.predict(data)
                prediction = [score.argmax(axis=1), score]
            else:
                prediction = [model.predict(data)]
        elif isinstance(model, xgboost.Booster):
            # XGBoost Booster
            from ..convert.xgboost._parse import _get_attributes
            from xgboost import DMatrix
            datax = DMatrix(data)
            model_dict = _get_attributes(model)
            if model_dict['objective'].startswith('binary'):
                score = model.predict(datax)
                prediction = [score > 0.5, numpy.vstack([1-score, score]).T]
            elif model_dict['objective'].startswith('multi:softprob'):
                score = model.predict(datax)
                prediction = [score.argmax(axis=1), score]
            elif model_dict['objective'].startswith('multi:softmax'):
                score = model.predict(datax, output_margin=True)
                prediction = [score.argmax(axis=1), score]
            else:
                prediction = [model.predict(datax)]
        elif hasattr(model, "predict_proba"):
            # Classifier
            if hasattr(model, 'get_params'):
                params = model.get_params()
                if 'objective' in params:
                    objective = params['objective']
                    if objective == "multi:softmax":
                        prediction = [model.predict(data)]
                    else:
                        prediction = [model.predict(data), model.predict_proba(data)]
                else:
                    prediction = [model.predict(data), model.predict_proba(data)]
            else:
                prediction = [model.predict(data), model.predict_proba(data)]
        elif hasattr(model, "predict_with_probabilities"):
            # Classifier that returns all in one go
            prediction = model.predict_with_probabilities(data)
        elif hasattr(model, "decision_function"):
            # Classifier without probabilities
            prediction = [model.predict(data), model.decision_function(data)]
        elif hasattr(model, "layers"):
            # Keras
            if len(model.input_names) != 1:
                raise NotImplemented("Only neural network with one input are supported")
            prediction = [model.predict(data)]
        else:
            # Regressor
            prediction = [model.predict(data)]
    elif hasattr(model, "transform"):
        prediction = model.transform(data)
    else:
        raise TypeError("Model has not predict or transform method: {0}".format(type(model)))

    runtime_test['expected'] = prediction

    names = []
    dest = os.path.join(folder, basename + ".expected.pkl")
    names.append(dest)
    with open(dest, "wb") as f:
        pickle.dump(prediction, f)

    dest = os.path.join(folder, basename + ".data.pkl")
    names.append(dest)
    with open(dest, "wb") as f:
        pickle.dump(data, f)

    dest = os.path.join(folder, basename + ".model.pkl")
    names.append(dest)
    with open(dest, "wb") as f:
        pickle.dump(model, f)

    if onnx is None:
        array = numpy.array(data)
        if inputs is None:
            inputs = [('input', FloatTensorType(list(array.shape)))]
        onnx, _ = convert_model(model, basename, inputs)

    dest = os.path.join(folder, basename + ".model.onnx")
    names.append(dest)
    with open(dest, "wb") as f:
        f.write(onnx.SerializeToString())

    runtime_test["onnx"] = dest

    # backend
    if backend is not None:
        if not isinstance(backend, list):
            backend = [backend]
        for b in backend:
            if not is_backend_enabled(b):
                continue
            if isinstance(allow_failure, str):
                allow = evaluate_condition(b, allow_failure)
            else:
                allow = allow_failure
            if allow is None:
                output = compare_backend(b, runtime_test, options=extract_options(basename),
                                         context=context, verbose=verbose)
            else:
                try:
                    output = compare_backend(b, runtime_test, options=extract_options(basename),
                                             context=context, verbose=verbose)
                except AssertionError as e:
                    if isinstance(allow, bool) and allow:
                        warnings.warn("Issue with '{0}' due to {1}".format(basename, e))
                        continue
                    else:
                        raise e
            if output is not None:
                dest = os.path.join(folder, basename + ".backend.{0}.pkl".format(b))
                names.append(dest)
                with open(dest, "wb") as f:
                    pickle.dump(output, f)

    return names


def convert_model(model, name, input_types, without_onnx_ml=False, **kwargs):
    """
    Runs the appropriate conversion method.

    :param model: model
    :return: *onnx* model
    """
    from sklearn.base import BaseEstimator
    if model.__class__.__name__.startswith("LGBM"):
        from onnxmltools.convert import convert_lightgbm
        model, prefix = convert_lightgbm(model, name, input_types, without_onnx_ml=without_onnx_ml, **kwargs), "LightGbm"
    elif model.__class__.__name__.startswith("XGB"):
        from onnxmltools.convert import convert_xgboost
        model, prefix = convert_xgboost(model, name, input_types, **kwargs), "XGB"
    elif model.__class__.__name__ == 'Booster':
        import lightgbm
        if isinstance(model, lightgbm.Booster):
            from onnxmltools.convert import convert_lightgbm
            model, prefix = convert_lightgbm(model, name, input_types, without_onnx_ml=without_onnx_ml, **kwargs), "LightGbm"
        else:
            raise RuntimeError("Unable to convert model of type '{0}'.".format(type(model)))
    elif model.__class__.__name__.startswith("CatBoost"):
        from onnxmltools.convert import convert_catboost
        model, prefix = convert_catboost(model, name, input_types, **kwargs), "CatBoost"
    elif isinstance(model, BaseEstimator):
        from onnxmltools.convert import convert_sklearn
        model, prefix = convert_sklearn(model, name, input_types, **kwargs), "Sklearn"
    else:
        from onnxmltools.convert import convert_coreml
        model, prefix = convert_coreml(model, name, input_types, **kwargs), "Cml"
    if model is None:
        raise RuntimeError("Unable to convert model of type '{0}'.".format(type(model)))
    return model, prefix


def dump_one_class_classification(model, suffix="", folder=None, allow_failure=None):
    """
    Trains and dumps a model for a One Class outlier problem.
    The function trains a model and calls
    :func:`dump_data_and_model`.

    :param model: any model following *scikit-learn* API
    :param suffix: added to filenames
    :param folder: where to save the file
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :return: output of :func:`dump_data_and_model`

    Every created filename will follow the pattern:
    ``<folder>/<prefix><task><classifier-name><suffix>.<data|expected|model|onnx>.<pkl|onnx>``.
    """
    X = [[0., 1.], [1., 1.], [2., 0.]]
    X = numpy.array(X, dtype=numpy.float32)
    y = [1, 1, 1]
    model.fit(X, y)
    model_onnx, prefix = convert_model(model, 'one_class', [('input', FloatTensorType([None, 2]))],
                                       target_opset=TARGET_OPSET)
    return dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure,
                               basename=prefix + "One" + model.__class__.__name__ + suffix)


def dump_binary_classification(model, suffix="", folder=None, allow_failure=None, verbose=False):
    """
    Trains and dumps a model for a binary classification problem.

    :param model: any model following *scikit-learn* API
    :param suffix: added to filenames
    :param folder: where to save the file
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :param verbose: prints more information when it fails
    :return: output of :func:`dump_data_and_model`

    Every created filename will follow the pattern:
    ``<folder>/<prefix><task><classifier-name><suffix>.<data|expected|model|onnx>.<pkl|onnx>``.
    """
    X = [[0, 1], [1, 1], [2, 0]]
    X = numpy.array(X, dtype=numpy.float32)
    y = [0, 1, 0]
    model.fit(X, y)
    model_onnx, prefix = convert_model(model, 'tree-based binary classifier', [('input', FloatTensorType([None, 2]))],
                                       target_opset=TARGET_OPSET)
    dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure,
                        basename=prefix + "Bin" + model.__class__.__name__ + suffix,
                        verbose=verbose)

def dump_multiple_classification(model, suffix="", folder=None, allow_failure=None):
    """
    Trains and dumps a model for a binary classification problem.

    :param model: any model following *scikit-learn* API
    :param suffix: added to filenames
    :param folder: where to save the file
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :return: output of :func:`dump_data_and_model`

    Every created filename will follow the pattern:
    ``<folder>/<prefix><task><classifier-name><suffix>.<data|expected|model|onnx>.<pkl|onnx>``.
    """
    X = [[0, 1], [1, 1], [2, 0], [0.5, 0.5], [1.1, 1.1], [2.1, 0.1]]
    X = numpy.array(X, dtype=numpy.float32)
    y = [0, 1, 2, 1, 1, 2]
    model.fit(X, y)
    model_onnx, prefix = convert_model(model, 'tree-based multi-output regressor', [('input', FloatTensorType([None, 2]))],
                                       target_opset=TARGET_OPSET)
    dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure,
                        basename=prefix + "Mcl" + model.__class__.__name__ + suffix)


def dump_multiple_regression(model, suffix="", folder=None, allow_failure=None):
    """
    Trains and dumps a model for a multi regression problem.

    :param model: any model following *scikit-learn* API
    :param suffix: added to filenames
    :param folder: where to save the file
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :return: output of :func:`dump_data_and_model`

    Every created filename will follow the pattern:
    ``<folder>/<prefix><task><classifier-name><suffix>.<data|expected|model|onnx>.<pkl|onnx>``.
    """
    X = [[0, 1], [1, 1], [2, 0]]
    X = numpy.array(X, dtype=numpy.float32)
    y = numpy.array([[100, 50], [100, 49], [100, 99]], dtype=numpy.float32)
    model.fit(X, y)
    model_onnx, prefix = convert_model(model, 'tree-based multi-output regressor', [('input', FloatTensorType([None, 2]))],
                                       target_opset=TARGET_OPSET)
    dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure,
                        basename=prefix + "MRg" + model.__class__.__name__ + suffix)


def dump_single_regression(model, suffix="", folder=None, allow_failure=None):
    """
    Trains and dumps a model for a regression problem.

    :param model: any model following *scikit-learn* API
    :param prefix: library name
    :param suffix: added to filenames
    :param folder: where to save the file
    :param allow_failure: None to raise an exception if comparison fails
        for the backends, otherwise a string which is then evaluated to check
        whether or not the test can fail, example:
        ``"pv.Version(onnx.__version__) < pv.Version('1.3.0')"``
    :return: output of :func:`dump_data_and_model`

    Every created filename will follow the pattern:
    ``<folder>/<prefix><task><classifier-name><suffix>.<data|expected|model|onnx>.<pkl|onnx>``.
    """
    X = [[0, 1], [1, 1], [2, 0]]
    X = numpy.array(X, dtype=numpy.float32)
    y = numpy.array([100, -10, 50], dtype=numpy.float32)
    model.fit(X, y)
    model_onnx, prefix = convert_model(model, 'tree-based regressor', [('input', FloatTensorType([None, 2]))],
                                       target_opset=TARGET_OPSET)
    dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure,
                        basename=prefix + "Reg" + model.__class__.__name__ + suffix)


def make_report_backend(folder):
    """
    Looks into a folder for dumped files after
    the unit tests.
    """
    res = {}
    files = os.listdir(folder)
    for name in files:
        if name.endswith(".expected.pkl"):
            model = name.split(".")[0]
            if model not in res:
                res[model] = {}
            res[model]["_tested"] = True
        elif '.backend.' in name:
            bk = name.split(".backend.")[-1].split(".")[0]
            model = name.split(".")[0]
            if model not in res:
                res[model] = {}
            res[model][bk] = True

    def dict_update(d, u):
        d.update(u)
        return d

    aslist = [dict_update(dict(_model=k), v) for k, v in res.items()]
    return aslist
