#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import tempfile
import unittest

import numpy as np

from pyspark.ml.linalg import Vectors
from pyspark.ml.regression import (
    AFTSurvivalRegression,
    AFTSurvivalRegressionModel,
    IsotonicRegression,
    IsotonicRegressionModel,
    LinearRegression,
    LinearRegressionModel,
    GeneralizedLinearRegression,
    GeneralizedLinearRegressionModel,
    GeneralizedLinearRegressionSummary,
    GeneralizedLinearRegressionTrainingSummary,
    LinearRegressionSummary,
    LinearRegressionTrainingSummary,
    FMRegressor,
    FMRegressionModel,
    DecisionTreeRegressor,
    DecisionTreeRegressionModel,
    RandomForestRegressor,
    RandomForestRegressionModel,
    GBTRegressor,
    GBTRegressionModel,
)
from pyspark.sql import is_remote
from pyspark.testing.sqlutils import ReusedSQLTestCase


class RegressionTestsMixin:
    def test_aft_survival(self):
        spark = self.spark
        df = spark.createDataFrame(
            [(1.0, Vectors.dense(1.0), 1.0), (1e-40, Vectors.sparse(1, [], []), 0.0)],
            ["label", "features", "censor"],
        )

        aft = AFTSurvivalRegression()
        aft.setMaxIter(1)
        self.assertEqual(aft.getMaxIter(), 1)

        model = aft.fit(df)
        self.assertEqual(aft.uid, model.uid)
        self.assertEqual(model.numFeatures, 1)
        self.assertTrue(np.allclose(model.intercept, 0.0, atol=1e-4), model.intercept)
        self.assertTrue(
            np.allclose(model.coefficients.toArray(), [0.0], atol=1e-4), model.coefficients
        )
        self.assertTrue(np.allclose(model.scale, 1.0, atol=1e-4), model.scale)

        vec = Vectors.dense(6.3)
        pred = model.predict(vec)
        self.assertEqual(pred, 1.0)
        pred = model.predictQuantiles(vec)
        self.assertTrue(
            np.allclose(
                pred,
                [
                    0.010050335853501444,
                    0.051293294387550536,
                    0.1053605156578263,
                    0.2876820724517809,
                    0.6931471805599453,
                    1.3862943611198906,
                    2.302585092994046,
                    2.9957322735539895,
                    4.60517018598809,
                ],
                atol=1e-4,
            ),
            pred,
        )

        output = model.transform(df)
        expected_cols = ["label", "features", "censor", "prediction"]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 2)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="aft_survival") as d:
            aft.write().overwrite().save(d)
            aft2 = AFTSurvivalRegression.load(d)
            self.assertEqual(str(aft), str(aft2))

            model.write().overwrite().save(d)
            model2 = AFTSurvivalRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_isotonic_regression(self):
        spark = self.spark
        df = spark.createDataFrame(
            [(1.0, Vectors.dense(1.0)), (0.0, Vectors.sparse(1, [], []))], ["label", "features"]
        )

        ir = IsotonicRegression(
            isotonic=True,
            featureIndex=0,
        )
        self.assertTrue(ir.getIsotonic())
        self.assertEqual(ir.getFeatureIndex(), 0)

        model = ir.fit(df)
        self.assertEqual(model.numFeatures, 1)
        self.assertTrue(
            np.allclose(model.boundaries.toArray(), [0.0, 1.0], atol=1e-4), model.boundaries
        )
        self.assertTrue(
            np.allclose(model.predictions.toArray(), [0.0, 1.0], atol=1e-4), model.predictions
        )

        pred = model.predict(1.0)
        self.assertTrue(np.allclose(pred, 1.0, atol=1e-4), pred)

        output = model.transform(df)
        expected_cols = ["label", "features", "prediction"]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 2)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="isotonic_regression") as d:
            ir.write().overwrite().save(d)
            ir2 = IsotonicRegression.load(d)
            self.assertEqual(str(ir), str(ir2))

            model.write().overwrite().save(d)
            model2 = IsotonicRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_linear_regression(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.5, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.7, 4.0, Vectors.dense(1.5, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )

        lr = LinearRegression(
            regParam=0.0,
            maxIter=2,
            solver="normal",
            weightCol="weight",
        )
        self.assertEqual(lr.getRegParam(), 0)
        self.assertEqual(lr.getMaxIter(), 2)
        self.assertEqual(lr.getSolver(), "normal")
        self.assertEqual(lr.getWeightCol(), "weight")

        model = lr.fit(df)
        self.assertEqual(lr.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)
        self.assertTrue(np.allclose(model.scale, 1.0, atol=1e-4))
        self.assertTrue(np.allclose(model.intercept, -0.35, atol=1e-4))
        self.assertTrue(np.allclose(model.coefficients, [0.65, 0.1125], atol=1e-4))

        output = model.transform(df)
        expected_cols = [
            "label",
            "weight",
            "features",
            "prediction",
        ]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        self.assertTrue(
            np.allclose(model.predict(Vectors.dense(0.0, 5.0)), 0.21249999999999963, atol=1e-4)
        )

        def check_summary():
            # Model summary
            summary = model.summary
            self.assertTrue(isinstance(summary, LinearRegressionSummary))
            self.assertTrue(isinstance(summary, LinearRegressionTrainingSummary))
            self.assertEqual(summary.predictions.columns, expected_cols)
            self.assertEqual(summary.predictions.count(), 4)
            self.assertEqual(summary.residuals.columns, ["residuals"])
            self.assertEqual(summary.residuals.count(), 4)

            self.assertEqual(summary.degreesOfFreedom, 1)
            self.assertEqual(summary.numInstances, 4)
            self.assertEqual(summary.objectiveHistory, [0.0])
            self.assertTrue(
                np.allclose(
                    summary.coefficientStandardErrors,
                    [1.2859821149611763, 0.6248749874975031, 3.1645497310044184],
                    atol=1e-4,
                )
            )
            self.assertTrue(
                np.allclose(
                    summary.devianceResiduals, [-0.7424621202458727, 0.7875000000000003], atol=1e-4
                )
            )
            self.assertTrue(
                np.allclose(
                    summary.pValues,
                    [0.7020630236843428, 0.8866003086182783, 0.9298746994547682],
                    atol=1e-4,
                )
            )
            self.assertTrue(
                np.allclose(
                    summary.tValues,
                    [0.5054502643838291, 0.1800360108036021, -0.11060025272186746],
                    atol=1e-4,
                )
            )
            self.assertTrue(np.allclose(summary.explainedVariance, 0.07997500000000031, atol=1e-4))
            self.assertTrue(np.allclose(summary.meanAbsoluteError, 0.4200000000000002, atol=1e-4))
            self.assertTrue(np.allclose(summary.meanSquaredError, 0.20212500000000005, atol=1e-4))
            self.assertTrue(
                np.allclose(summary.rootMeanSquaredError, 0.44958314025327956, atol=1e-4)
            )
            self.assertTrue(np.allclose(summary.r2, 0.4427212572373862, atol=1e-4))
            self.assertTrue(np.allclose(summary.r2adj, -0.6718362282878414, atol=1e-4))

        check_summary()
        if is_remote():
            self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
            check_summary()

        summary2 = model.evaluate(df)
        self.assertTrue(isinstance(summary2, LinearRegressionSummary))
        self.assertFalse(isinstance(summary2, LinearRegressionTrainingSummary))
        self.assertEqual(summary2.predictions.columns, expected_cols)
        self.assertEqual(summary2.predictions.count(), 4)
        self.assertEqual(summary2.residuals.columns, ["residuals"])
        self.assertEqual(summary2.residuals.count(), 4)

        self.assertEqual(summary2.degreesOfFreedom, 1)
        self.assertEqual(summary2.numInstances, 4)
        self.assertTrue(
            np.allclose(
                summary2.devianceResiduals, [-0.7424621202458727, 0.7875000000000003], atol=1e-4
            )
        )
        self.assertTrue(np.allclose(summary2.explainedVariance, 0.07997500000000031, atol=1e-4))
        self.assertTrue(np.allclose(summary2.meanAbsoluteError, 0.4200000000000002, atol=1e-4))
        self.assertTrue(np.allclose(summary2.meanSquaredError, 0.20212500000000005, atol=1e-4))
        self.assertTrue(np.allclose(summary2.rootMeanSquaredError, 0.44958314025327956, atol=1e-4))
        self.assertTrue(np.allclose(summary2.r2, 0.4427212572373862, atol=1e-4))
        self.assertTrue(np.allclose(summary2.r2adj, -0.6718362282878414, atol=1e-4))

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="linear_regression") as d:
            lr.write().overwrite().save(d)
            lr2 = LinearRegression.load(d)
            self.assertEqual(str(lr), str(lr2))

            model.write().overwrite().save(d)
            model2 = LinearRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_generalized_linear_regression(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense(0.0, 0.0)),
                    (2, 1.0, Vectors.dense(1.0, 2.0)),
                    (3, 2.0, Vectors.dense(0.0, 0.0)),
                    (4, 2.0, Vectors.dense(1.0, 1.0)),
                ],
                ["index", "label", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("label", "features")
        )

        glr = GeneralizedLinearRegression(
            family="gaussian",
            link="identity",
            linkPredictionCol="p",
        )
        glr.setRegParam(0.1)
        glr.setMaxIter(1)
        self.assertEqual(glr.getFamily(), "gaussian")
        self.assertEqual(glr.getLink(), "identity")
        self.assertEqual(glr.getLinkPredictionCol(), "p")
        self.assertEqual(glr.getRegParam(), 0.1)
        self.assertEqual(glr.getMaxIter(), 1)

        model = glr.fit(df)
        self.assertTrue(np.allclose(model.intercept, 1.543859649122807, atol=1e-4), model.intercept)
        self.assertTrue(
            np.allclose(model.coefficients.toArray(), [0.43859649, -0.35087719], atol=1e-4),
            model.coefficients,
        )
        self.assertEqual(model.numFeatures, 2)

        vec = Vectors.dense(1.0, 2.0)
        pred = model.predict(vec)
        self.assertTrue(np.allclose(pred, 1.280701754385965, atol=1e-4), pred)

        expected_cols = ["label", "features", "p", "prediction"]

        output = model.transform(df)
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        def check_summary():
            # Model summary
            self.assertTrue(model.hasSummary)

            summary = model.summary
            self.assertIsInstance(summary, GeneralizedLinearRegressionSummary)
            self.assertIsInstance(summary, GeneralizedLinearRegressionTrainingSummary)
            self.assertEqual(summary.numIterations, 1)
            self.assertEqual(summary.numInstances, 4)
            self.assertEqual(summary.rank, 3)
            self.assertTrue(
                np.allclose(
                    summary.tValues,
                    [0.3725037662281711, -0.49418209022924164, 2.6589353685797654],
                    atol=1e-4,
                ),
                summary.tValues,
            )
            self.assertTrue(
                np.allclose(
                    summary.pValues,
                    [0.7729938686180984, 0.707802691825973, 0.22900885781807023],
                    atol=1e-4,
                ),
                summary.pValues,
            )
            self.assertEqual(summary.predictions.columns, expected_cols)
            self.assertEqual(summary.predictions.count(), 4)
            self.assertEqual(summary.residuals().columns, ["devianceResiduals"])
            self.assertEqual(summary.residuals().count(), 4)

        check_summary()
        if is_remote():
            self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
            check_summary()

        summary = model.summary
        summary2 = model.evaluate(df)
        self.assertIsInstance(summary2, GeneralizedLinearRegressionSummary)
        self.assertNotIsInstance(summary2, GeneralizedLinearRegressionTrainingSummary)
        self.assertEqual(summary2.numInstances, 4)
        self.assertEqual(summary2.rank, 3)
        self.assertEqual(summary.predictions.columns, expected_cols)
        self.assertEqual(summary.predictions.count(), 4)
        self.assertEqual(summary2.residuals().columns, ["devianceResiduals"])
        self.assertEqual(summary2.residuals().count(), 4)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="generalized_linear_regression") as d:
            glr.write().overwrite().save(d)
            glr2 = GeneralizedLinearRegression.load(d)
            self.assertEqual(str(glr), str(glr2))

            model.write().overwrite().save(d)
            model2 = GeneralizedLinearRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_factorization_machine(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense(0.0, 0.0)),
                    (2, 1.0, Vectors.dense(1.0, 2.0)),
                    (3, 2.0, Vectors.dense(0.0, 0.0)),
                    (4, 2.0, Vectors.dense(1.0, 1.0)),
                ],
                ["index", "label", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("label", "features")
        )

        fm = FMRegressor(factorSize=2, maxIter=1, regParam=1.0, seed=1)
        self.assertEqual(fm.getFactorSize(), 2)
        self.assertEqual(fm.getMaxIter(), 1)
        self.assertEqual(fm.getRegParam(), 1.0)
        self.assertEqual(fm.getSeed(), 1)

        model = fm.fit(df)
        self.assertEqual(fm.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)
        self.assertTrue(
            np.allclose(model.intercept, 0.9999999966668874, atol=1e-4), model.intercept
        )
        self.assertTrue(
            np.allclose(
                model.linear.toArray(), [0.9999999933342161, 0.9999999950008276], atol=1e-4
            ),
            model.linear,
        )
        self.assertTrue(
            np.allclose(
                model.factors.toArray(),
                [[-0.99999954, -0.9999992], [0.99999968, -0.99999918]],
                atol=1e-4,
            ),
            model.factors,
        )

        vec = Vectors.dense(0.0, 5.0)
        pred = model.predict(vec)
        self.assertTrue(np.allclose(pred, 5.999999971671025, atol=1e-4), pred)

        output = model.transform(df)
        expected_cols = ["label", "features", "prediction"]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="factorization_machine") as d:
            fm.write().overwrite().save(d)
            fm2 = FMRegressor.load(d)
            self.assertEqual(str(fm), str(fm2))

            model.write().overwrite().save(d)
            model2 = FMRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))

    def test_decision_tree_regressor(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.5, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.7, 4.0, Vectors.dense(1.5, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )

        dt = DecisionTreeRegressor(
            maxDepth=2,
            labelCol="label",
            leafCol="leaf",
            seed=1,
        )
        self.assertEqual(dt.getMaxDepth(), 2)
        self.assertEqual(dt.getSeed(), 1)
        self.assertEqual(dt.getLabelCol(), "label")
        self.assertEqual(dt.getLeafCol(), "leaf")

        model = dt.fit(df)
        self.assertEqual(dt.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)
        self.assertEqual(model.depth, 2)
        self.assertEqual(model.numNodes, 5)

        featureImportances = model.featureImportances
        self.assertTrue(
            np.allclose(featureImportances, [0.5756, 0.4244], atol=1e-4),
            featureImportances,
        )

        debugString = model.toDebugString
        self.assertTrue("depth=2, numNodes=5, numFeatures=2" in debugString, debugString)
        self.assertTrue("If (feature 0 <= 1.75)" in debugString, debugString)

        vec = Vectors.dense(0.0, 5.0)
        self.assertTrue(np.allclose(model.predict(vec), 0.85, atol=1e-4))
        self.assertEqual(model.predictLeaf(vec), 1.0)

        output = model.transform(df)
        expected_cols = [
            "label",
            "weight",
            "features",
            "prediction",
            "leaf",
        ]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        # Model save & load
        with tempfile.TemporaryDirectory(prefix="decision_tree_regression") as d:
            dt.write().overwrite().save(d)
            dt2 = DecisionTreeRegressor.load(d)
            self.assertEqual(str(dt), str(dt2))

            model.write().overwrite().save(d)
            model2 = DecisionTreeRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.toDebugString, model2.toDebugString)

    def test_gbt_regressor(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.5, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.7, 4.0, Vectors.dense(1.5, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )

        gbt = GBTRegressor(
            maxIter=3,
            maxDepth=2,
            labelCol="label",
            leafCol="leaf",
            seed=1,
        )
        self.assertEqual(gbt.getMaxIter(), 3)
        self.assertEqual(gbt.getMaxDepth(), 2)
        self.assertEqual(gbt.getSeed(), 1)
        self.assertEqual(gbt.getLabelCol(), "label")
        self.assertEqual(gbt.getLeafCol(), "leaf")

        model = gbt.fit(df)
        self.assertEqual(gbt.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)
        self.assertEqual(model.treeWeights, [1.0, 0.1, 0.1])
        self.assertEqual(model.totalNumNodes, 15)

        featureImportances = model.featureImportances
        self.assertTrue(
            np.allclose(featureImportances, [0.5944156994359766, 0.4055843005640234], atol=1e-4),
            featureImportances,
        )

        debugString = model.toDebugString
        self.assertTrue("numTrees=3, numFeatures=2" in debugString, debugString)
        self.assertTrue("If (feature 0 <= 1.75)" in debugString, debugString)

        vec = Vectors.dense(0.0, 5.0)
        self.assertTrue(np.allclose(model.predict(vec), 0.904, atol=1e-4))
        self.assertEqual(model.predictLeaf(vec), Vectors.dense(1.0, 0.0, 0.0))

        # GBT-specific method: evaluateEachIteration
        self.assertTrue(
            np.allclose(
                model.evaluateEachIteration(df, "squared"),
                [0.011250000000000003, 0.0072, 0.0046079999999999975],
                atol=1e-4,
            )
        )
        self.assertTrue(
            np.allclose(
                model.evaluateEachIteration(df, "absolute"),
                [0.07500000000000007, 0.06000000000000006, 0.048000000000000057],
                atol=1e-4,
            )
        )

        output = model.transform(df)
        expected_cols = [
            "label",
            "weight",
            "features",
            "prediction",
            "leaf",
        ]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        trees = model.trees
        self.assertEqual(len(trees), 3)
        for tree in trees:
            self.assertIsInstance(tree, DecisionTreeRegressionModel)
            self.assertTrue(tree.predict(vec) > -10)
            self.assertEqual(tree.transform(df).count(), 4)
            self.assertEqual(
                tree.transform(df).columns,
                ["label", "weight", "features", "prediction", "leaf"],
            )

        # save & load
        with tempfile.TemporaryDirectory(prefix="gbt_regression") as d:
            gbt.write().overwrite().save(d)
            gbt2 = GBTRegressor.load(d)
            self.assertEqual(str(gbt), str(gbt2))

            model.write().overwrite().save(d)
            model2 = GBTRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.toDebugString, model2.toDebugString)

    def test_random_forest_regressor(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.5, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.7, 4.0, Vectors.dense(1.5, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )

        rf = RandomForestRegressor(
            numTrees=3,
            maxDepth=2,
            labelCol="label",
            leafCol="leaf",
            seed=1,
        )
        self.assertEqual(rf.getNumTrees(), 3)
        self.assertEqual(rf.getMaxDepth(), 2)
        self.assertEqual(rf.getSeed(), 1)
        self.assertEqual(rf.getLabelCol(), "label")
        self.assertEqual(rf.getLeafCol(), "leaf")

        model = rf.fit(df)
        self.assertEqual(rf.uid, model.uid)
        self.assertEqual(model.numFeatures, 2)
        self.assertEqual(model.treeWeights, [1.0, 1.0, 1.0])
        self.assertEqual(model.totalNumNodes, 11)

        featureImportances = model.featureImportances
        self.assertTrue(
            np.allclose(featureImportances, [0.5615222294986538, 0.43847777050134623], atol=1e-4),
            featureImportances,
        )

        debugString = model.toDebugString
        self.assertTrue("numTrees=3, numFeatures=2" in debugString, debugString)
        self.assertTrue("If (feature 0 <= 1.75)" in debugString, debugString)

        vec = Vectors.dense(0.0, 5.0)
        self.assertTrue(np.allclose(model.predict(vec), 0.6166666666666667, atol=1e-4))
        self.assertEqual(model.predictLeaf(vec), Vectors.dense(1.0, 0.0, 1.0))

        output = model.transform(df)
        expected_cols = [
            "label",
            "weight",
            "features",
            "prediction",
            "leaf",
        ]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 4)

        trees = model.trees
        self.assertEqual(len(trees), 3)
        for tree in trees:
            self.assertIsInstance(tree, DecisionTreeRegressionModel)
            self.assertTrue(tree.predict(vec) > -10)
            self.assertEqual(tree.transform(df).count(), 4)
            self.assertEqual(
                tree.transform(df).columns,
                ["label", "weight", "features", "prediction", "leaf"],
            )

        with tempfile.TemporaryDirectory(prefix="random_forest_regression") as d:
            rf.write().overwrite().save(d)
            rf2 = RandomForestRegressor.load(d)
            self.assertEqual(str(rf), str(rf2))

            model.write().overwrite().save(d)
            model2 = RandomForestRegressionModel.load(d)
            self.assertEqual(str(model), str(model2))
            self.assertEqual(model.toDebugString, model2.toDebugString)


class RegressionTests(RegressionTestsMixin, ReusedSQLTestCase):
    pass


if __name__ == "__main__":
    from pyspark.ml.tests.test_regression import *  # noqa: F401,F403

    try:
        import xmlrunner  # type: ignore[import]

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)