# SPDX-License-Identifier: Apache-2.0

import inspect
import unittest
import os
import numpy
import pandas
from pyspark.ml import PipelineModel
from onnxmltools import convert_sparkml
from onnxmltools.convert.sparkml import buildInitialTypesSimple, buildInputDictSimple
from onnxmltools.utils.utils_backend_onnxruntime import (
    run_with_runtime,
    _compare_expected,
)
from tests.sparkml import SparkMlTestCase


class RPipeline(SparkMlTestCase):
    def test_sparkml_r_pipeline(self):
        # add additional jar files before creating SparkSession
        this_script_dir = os.path.dirname(
            os.path.abspath(inspect.getfile(inspect.currentframe()))
        )
        input_path = os.path.join(this_script_dir, "data", "iris.csv")
        data = (
            self.spark.read.format("csv")
            .options(header="true", inferschema="true")
            .load(input_path)
            .drop("_index_")
        )

        # read the model from disk
        pipeline_path = os.path.join(this_script_dir, "mlpmodel")
        model = PipelineModel.load(path=pipeline_path)

        # create Onnx model
        model_onnx = convert_sparkml(
            model,
            "Sparkml R Pipeline",
            buildInitialTypesSimple(data),
            spark_session=self.spark,
        )
        # save Onnx model for runtime usage
        if model_onnx is None:
            raise AssertionError("Failed to create the onnx model")
        model_path = os.path.join(
            this_script_dir, "tests_dump", "r_pipeline_model.onnx"
        )
        with open(model_path, "wb") as f:
            f.write(model_onnx.SerializeToString())

        data_np = buildInputDictSimple(data)
        # run the model in Spark
        spark_prediction = model.transform(data)
        # run the model in onnx runtime
        output, session = run_with_runtime(data_np, model_path)

        # compare results
        expected = [
            spark_prediction.toPandas().label.values.astype(numpy.float32),
            spark_prediction.toPandas().prediction.values.astype(numpy.float32),
            spark_prediction.toPandas()
            .probability.apply(lambda x: pandas.Series(x.toArray()))
            .values.astype(numpy.float32),
        ]
        _compare_expected(
            expected, output, session, model_path, decimal=5, onnx_shape=None
        )


if __name__ == "__main__":
    unittest.main()
