# SPDX-License-Identifier: Apache-2.0

import sys
import unittest
import numpy
from pyspark.ml.feature import Imputer
from onnx.defs import onnx_opset_version
from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER
from onnxmltools import convert_sparkml
from onnxmltools.convert.common.data_types import FloatTensorType
from tests.sparkml.sparkml_test_utils import (
    save_data_models,
    run_onnx_model,
    compare_results,
)
from tests.sparkml import SparkMlTestCase


TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version())


## For some reason during the spark bring up and shutdown
## something happens causing Imputer
## tests to fail. For that you need to run each test here individually
## for now these will be commented out so as not to break the build
##      AttributeError: 'NoneType' object has no attribute
# 'setCallSite' on model.surrogateDF
## Therefore we leave these tests out for now until a newere
## version of pyspark is availabe that address this issue
class TestSparkmlImputer(SparkMlTestCase):
    @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7")
    def test_imputer_single(self):
        self._imputer_test_single()

    @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7")
    def test_imputer_multi(self):
        self._imputer_test_multi()

    def _imputer_test_multi(self):
        data = self.spark.createDataFrame(
            [
                (1.0, float("nan")),
                (2.0, float("nan")),
                (float("nan"), 3.0),
                (4.0, 4.0),
                (5.0, 5.0),
            ],
            ["a", "b"],
        )
        imputer = Imputer(inputCols=["a", "b"], outputCols=["out_a", "out_b"])
        model = imputer.fit(data)

        # the input name should match the inputCols above
        model_onnx = convert_sparkml(
            model,
            "Sparkml Imputer Multi Input",
            [("a", FloatTensorType([None, 1])), ("b", FloatTensorType([None, 1]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx is not None)

        # run the model
        predicted = model.transform(data)

        expected = {
            "out_a": predicted.select("out_a").toPandas().values.astype(numpy.int64),
            "out_b": predicted.select("out_b").toPandas().values.astype(numpy.int64),
        }

        data_np = data.toPandas().values.astype(numpy.float32)
        data_np = {"a": data_np[:, :1], "b": data_np[:, 1:]}
        paths = save_data_models(
            data_np, expected, model, model_onnx, basename="SparkmlImputerMulti"
        )
        onnx_model_path = paths[-1]
        output_names = ["out_a", "out_b"]
        output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path)
        actual_output = dict(zip(output_names, output))
        compare_results(expected, actual_output, decimal=5)

    def _imputer_test_single(self):
        data = self.spark.createDataFrame(
            [
                (1.0, float("nan")),
                (2.0, float("nan")),
                (float("nan"), 3.0),
                (4.0, 4.0),
                (5.0, 5.0),
            ],
            ["a", "b"],
        )
        imputer = Imputer(inputCols=["a"], outputCols=["out_a"])
        model = imputer.fit(data)

        # the input name should match the inputCols above
        model_onnx = convert_sparkml(
            model,
            "Sparkml Imputer",
            [("a", FloatTensorType([None, 1]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx is not None)

        # run the model
        predicted = model.transform(data)
        expected = predicted.select("out_a").toPandas().values.astype(numpy.float32)
        data_np = data.toPandas().a.values.astype(numpy.float32)
        data_np = data_np.reshape((-1, 1))
        paths = save_data_models(
            data_np, expected, model, model_onnx, basename="SparkmlImputerSingle"
        )
        onnx_model_path = paths[-1]
        output, output_shapes = run_onnx_model(["out_a"], data_np, onnx_model_path)
        compare_results(expected, output, decimal=5)


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