"""
Tests scikit-learn's polynomial features converter.
"""
import unittest
from distutils.version import StrictVersion
import numpy as np
import onnx
from sklearn.preprocessing import PolynomialFeatures
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType, Int64TensorType
from test_utils import dump_data_and_model


class TestSklearnPolynomialFeatures(unittest.TestCase):

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_float_degree_2(self):
        X = np.array([[1.2, 3.2, 1.3, -5.6], [4.3, -3.2, 5.7, 1.0],
                      [0, 3.2, 4.7, -8.9]])
        model = PolynomialFeatures(degree=2).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", FloatTensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.float32),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesFloatDegree2",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_int_degree_2(self):
        X = np.array([
            [1, 3, 4, 0],
            [2, 3, 4, 1],
            [1, -4, 3, 7],
            [3, 10, -9, 5],
            [1, 0, 10, 5],
        ])
        model = PolynomialFeatures(degree=2).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", Int64TensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.int64),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesIntDegree2",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_float_degree_3(self):
        X = np.array([[1.2, 3.2, 1.2], [4.3, 3.2, 4.5], [3.2, 4.7, 1.1]])
        model = PolynomialFeatures(degree=3).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", FloatTensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.float32),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesFloatDegree3",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_int_degree_3(self):
        X = np.array([
            [1, 3, 33],
            [4, 1, -11],
            [3, 7, -3],
            [3, 5, 4],
            [1, 0, 3],
            [5, 4, 9],
        ])
        model = PolynomialFeatures(degree=3).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", Int64TensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.int64),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesIntDegree3",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_float_degree_4(self):
        X = np.array([[1.2, 3.2, 3.1, 1.3], [4.3, 3.2, 0.5, 1.3],
                      [3.2, 4.7, 5.4, 7.1]])
        model = PolynomialFeatures(degree=4).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", FloatTensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.float32),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesFloatDegree4-Dec4",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )

    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion("1.4.0"),
                     reason="ConstantOfShape not available")
    def test_model_polynomial_features_int_degree_4(self):
        X = np.array([[1, 3, 4, 1], [3, 7, 3, 5], [1, 0, 5, 4]])
        model = PolynomialFeatures(degree=4).fit(X)
        model_onnx = convert_sklearn(
            model,
            "scikit-learn polynomial features",
            [("input", Int64TensorType([None, X.shape[1]]))],
        )
        self.assertTrue(model_onnx is not None)
        dump_data_and_model(
            X.astype(np.int64),
            model,
            model_onnx,
            basename="SklearnPolynomialFeaturesIntDegree4",
            allow_failure="StrictVersion(onnxruntime.__version__)"
                          " <= StrictVersion('0.2.1')",
        )


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