#
# 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, SparseVector
from pyspark.ml.clustering import (
    KMeans,
    KMeansModel,
    KMeansSummary,
    BisectingKMeans,
    BisectingKMeansModel,
    BisectingKMeansSummary,
    GaussianMixture,
    GaussianMixtureModel,
    GaussianMixtureSummary,
    LDA,
    LDAModel,
    LocalLDAModel,
    DistributedLDAModel,
    PowerIterationClustering,
)
from pyspark.testing.sqlutils import ReusedSQLTestCase


class ClusteringTestsMixin:
    def test_kmeans(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([-0.1, -0.05])),
                    (2, 2.0, Vectors.dense([-0.01, -0.1])),
                    (3, 3.0, Vectors.dense([0.9, 0.8])),
                    (4, 1.0, Vectors.dense([0.75, 0.935])),
                    (5, 1.0, Vectors.dense([-0.83, -0.68])),
                    (6, 1.0, Vectors.dense([-0.91, -0.76])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("weight", "features")
        )

        km = KMeans(
            k=2,
            maxIter=2,
            weightCol="weight",
        )
        self.assertEqual(km.getK(), 2)
        self.assertEqual(km.getMaxIter(), 2)
        self.assertEqual(km.getWeightCol(), "weight")

        model = km.fit(df)
        self.assertEqual(km.uid, model.uid)

        centers = model.clusterCenters()
        self.assertEqual(len(centers), 2)
        self.assertTrue(np.allclose(centers[0], [-0.372, -0.338], atol=1e-3), centers[0])
        self.assertTrue(np.allclose(centers[1], [0.8625, 0.83375], atol=1e-3), centers[1])

        # TODO: support KMeansModel.numFeatures in Python
        # self.assertEqual(model.numFeatures, 2)

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

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

        # Model summary
        self.assertTrue(model.hasSummary)
        summary = model.summary
        self.assertTrue(isinstance(summary, KMeansSummary))
        self.assertEqual(summary.k, 2)
        self.assertEqual(summary.numIter, 2)
        self.assertEqual(summary.clusterSizes, [4, 2])
        self.assertTrue(np.allclose(summary.trainingCost, 1.35710375, atol=1e-4))

        self.assertEqual(summary.featuresCol, "features")
        self.assertEqual(summary.predictionCol, "prediction")

        self.assertEqual(summary.cluster.columns, ["prediction"])
        self.assertEqual(summary.cluster.count(), 6)

        self.assertEqual(summary.predictions.columns, expected_cols)
        self.assertEqual(summary.predictions.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="kmeans_model") as d:
            km.write().overwrite().save(d)
            km2 = KMeans.load(d)
            self.assertEqual(str(km), str(km2))

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

    def test_bisecting_kmeans(self):
        df = (
            self.spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([-0.1, -0.05])),
                    (2, 2.0, Vectors.dense([-0.01, -0.1])),
                    (3, 3.0, Vectors.dense([0.9, 0.8])),
                    (4, 1.0, Vectors.dense([0.75, 0.935])),
                    (5, 1.0, Vectors.dense([-0.83, -0.68])),
                    (6, 1.0, Vectors.dense([-0.91, -0.76])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("weight", "features")
        )

        bkm = BisectingKMeans(
            k=2,
            maxIter=2,
            minDivisibleClusterSize=1.0,
            weightCol="weight",
        )
        self.assertEqual(bkm.getK(), 2)
        self.assertEqual(bkm.getMaxIter(), 2)
        self.assertEqual(bkm.getMinDivisibleClusterSize(), 1.0)
        self.assertEqual(bkm.getWeightCol(), "weight")

        model = bkm.fit(df)
        self.assertEqual(bkm.uid, model.uid)

        centers = model.clusterCenters()
        self.assertEqual(len(centers), 2)
        self.assertTrue(np.allclose(centers[0], [-0.372, -0.338], atol=1e-3), centers[0])
        self.assertTrue(np.allclose(centers[1], [0.8625, 0.83375], atol=1e-3), centers[1])

        # TODO: support KMeansModel.numFeatures in Python
        # self.assertEqual(model.numFeatures, 2)

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

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

        # BisectingKMeans-specific method: computeCost
        self.assertTrue(np.allclose(model.computeCost(df), 1.164325125, atol=1e-4))

        # Model summary
        self.assertTrue(model.hasSummary)
        summary = model.summary
        self.assertTrue(isinstance(summary, BisectingKMeansSummary))
        self.assertEqual(summary.k, 2)
        self.assertEqual(summary.numIter, 2)
        self.assertEqual(summary.clusterSizes, [4, 2])
        self.assertTrue(np.allclose(summary.trainingCost, 1.3571037499999998, atol=1e-4))

        self.assertEqual(summary.featuresCol, "features")
        self.assertEqual(summary.predictionCol, "prediction")

        self.assertEqual(summary.cluster.columns, ["prediction"])
        self.assertEqual(summary.cluster.count(), 6)

        self.assertEqual(summary.predictions.columns, expected_cols)
        self.assertEqual(summary.predictions.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="bisecting_kmeans") as d:
            bkm.write().overwrite().save(d)
            bkm2 = BisectingKMeans.load(d)
            self.assertEqual(str(bkm), str(bkm2))

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

    def test_gaussian_mixture(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    (1, 1.0, Vectors.dense([-0.1, -0.05])),
                    (2, 2.0, Vectors.dense([-0.01, -0.1])),
                    (3, 3.0, Vectors.dense([0.9, 0.8])),
                    (4, 1.0, Vectors.dense([0.75, 0.935])),
                    (5, 1.0, Vectors.dense([-0.83, -0.68])),
                    (6, 1.0, Vectors.dense([-0.91, -0.76])),
                ],
                ["index", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("index")
            .select("weight", "features")
        )

        gmm = GaussianMixture(
            k=2,
            maxIter=2,
            weightCol="weight",
            seed=1,
        )
        self.assertEqual(gmm.getK(), 2)
        self.assertEqual(gmm.getMaxIter(), 2)
        self.assertEqual(gmm.getWeightCol(), "weight")
        self.assertEqual(gmm.getSeed(), 1)

        model = gmm.fit(df)
        self.assertEqual(gmm.uid, model.uid)
        # TODO: support GMM.numFeatures in Python
        # self.assertEqual(model.numFeatures, 2)
        self.assertEqual(len(model.weights), 2)
        self.assertTrue(
            np.allclose(model.weights, [0.541014115744985, 0.4589858842550149], atol=1e-4),
            model.weights,
        )

        # check the gaussians
        gaussians = model.gaussians
        self.assertEqual(len(gaussians), 2)
        self.assertTrue(
            np.allclose(
                gaussians[0].mean.toArray(),
                [0.28586084899633746, 0.28513455726904297],
                atol=1e-4,
            ),
            gaussians[0].mean,
        )
        self.assertTrue(
            np.allclose(
                gaussians[0].cov.toArray(),
                [[0.41732752, 0.38378601], [0.38378601, 0.36454957]],
                atol=1e-4,
            ),
            gaussians[0].cov,
        )
        self.assertTrue(
            np.allclose(
                gaussians[1].mean.toArray(),
                [0.04795771063097124, 0.06212817950777127],
                atol=1e-4,
            ),
            gaussians[1].mean,
        )
        self.assertTrue(
            np.allclose(
                gaussians[1].cov.toArray(),
                [[0.50359595, 0.44696663], [0.44696663, 0.40424231]],
                atol=1e-4,
            ),
            gaussians[1].cov,
        )

        self.assertEqual(model.gaussiansDF.columns, ["mean", "cov"])
        self.assertEqual(model.gaussiansDF.count(), 2)

        vec = Vectors.dense(0.0, 5.0)
        pred = model.predict(vec)
        self.assertTrue(np.allclose(pred, 0, atol=1e-4), pred)
        pred = model.predictProbability(vec)
        self.assertTrue(np.allclose(pred.toArray(), [0.5, 0.5], atol=1e-4), pred)

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

        # Model summary
        self.assertTrue(model.hasSummary)
        summary = model.summary
        self.assertTrue(isinstance(summary, GaussianMixtureSummary))
        self.assertEqual(summary.k, 2)
        self.assertEqual(summary.numIter, 2)
        self.assertEqual(len(summary.clusterSizes), 2)
        self.assertEqual(summary.clusterSizes, [3, 3])
        ll = summary.logLikelihood
        self.assertTrue(ll < 0, ll)
        self.assertTrue(np.allclose(ll, -1.311264553744033, atol=1e-4), ll)

        self.assertEqual(summary.featuresCol, "features")
        self.assertEqual(summary.predictionCol, "prediction")
        self.assertEqual(summary.probabilityCol, "probability")

        self.assertEqual(summary.cluster.columns, ["prediction"])
        self.assertEqual(summary.cluster.count(), 6)

        self.assertEqual(summary.predictions.columns, expected_cols)
        self.assertEqual(summary.predictions.count(), 6)

        self.assertEqual(summary.probability.columns, ["probability"])
        self.assertEqual(summary.predictions.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="gaussian_mixture") as d:
            gmm.write().overwrite().save(d)
            gmm2 = GaussianMixture.load(d)
            self.assertEqual(str(gmm), str(gmm2))

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

    def test_local_lda(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    [1, Vectors.dense([0.0, 1.0])],
                    [2, SparseVector(2, {0: 1.0})],
                ],
                ["id", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("id")
        )

        lda = LDA(k=2, optimizer="online", seed=1)
        lda.setMaxIter(1)
        self.assertEqual(lda.getK(), 2)
        self.assertEqual(lda.getOptimizer(), "online")
        self.assertEqual(lda.getMaxIter(), 1)
        self.assertEqual(lda.getSeed(), 1)

        model = lda.fit(df)
        self.assertEqual(lda.uid, model.uid)
        self.assertIsInstance(model, LDAModel)
        self.assertIsInstance(model, LocalLDAModel)
        self.assertNotIsInstance(model, DistributedLDAModel)
        self.assertFalse(model.isDistributed())

        dc = model.estimatedDocConcentration()
        self.assertTrue(np.allclose(dc.toArray(), [0.5, 0.5], atol=1e-4), dc)
        topics = model.topicsMatrix()
        self.assertTrue(
            np.allclose(
                topics.toArray(), [[1.20296728, 1.15740442], [0.99357675, 1.02993164]], atol=1e-4
            ),
            topics,
        )

        ll = model.logLikelihood(df)
        self.assertTrue(np.allclose(ll, -3.2125122434040088, atol=1e-4), ll)
        lp = model.logPerplexity(df)
        self.assertTrue(np.allclose(lp, 1.6062561217020044, atol=1e-4), lp)
        dt = model.describeTopics()
        self.assertEqual(dt.columns, ["topic", "termIndices", "termWeights"])
        self.assertEqual(dt.count(), 2)

        # LocalLDAModel specific methods
        self.assertEqual(model.vocabSize(), 2)

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

        # save & load
        with tempfile.TemporaryDirectory(prefix="local_lda") as d:
            lda.write().overwrite().save(d)
            lda2 = LDA.load(d)
            self.assertEqual(str(lda), str(lda2))

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

    def test_distributed_lda(self):
        spark = self.spark
        df = (
            spark.createDataFrame(
                [
                    [1, Vectors.dense([0.0, 1.0])],
                    [2, SparseVector(2, {0: 1.0})],
                ],
                ["id", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("id")
        )

        lda = LDA(k=2, optimizer="em", seed=1)
        lda.setMaxIter(1)

        self.assertEqual(lda.getK(), 2)
        self.assertEqual(lda.getOptimizer(), "em")
        self.assertEqual(lda.getMaxIter(), 1)
        self.assertEqual(lda.getSeed(), 1)

        model = lda.fit(df)
        self.assertEqual(lda.uid, model.uid)
        self.assertIsInstance(model, LDAModel)
        self.assertNotIsInstance(model, LocalLDAModel)
        self.assertIsInstance(model, DistributedLDAModel)
        self.assertTrue(model.isDistributed())
        self.assertEqual(model.vocabSize(), 2)

        dc = model.estimatedDocConcentration()
        self.assertTrue(np.allclose(dc.toArray(), [26.0, 26.0], atol=1e-4), dc)
        topics = model.topicsMatrix()
        self.assertTrue(
            np.allclose(
                topics.toArray(), [[0.39149926, 0.60850074], [0.60991237, 0.39008763]], atol=1e-4
            ),
            topics,
        )

        ll = model.logLikelihood(df)
        self.assertTrue(np.allclose(ll, -3.719138517085772, atol=1e-4), ll)
        lp = model.logPerplexity(df)
        self.assertTrue(np.allclose(lp, 1.859569258542886, atol=1e-4), lp)

        dt = model.describeTopics()
        self.assertEqual(dt.columns, ["topic", "termIndices", "termWeights"])
        self.assertEqual(dt.count(), 2)

        # DistributedLDAModel specific methods
        ll = model.trainingLogLikelihood()
        self.assertTrue(np.allclose(ll, -1.3847360462201639, atol=1e-4), ll)
        lp = model.logPrior()
        self.assertTrue(np.allclose(lp, -69.59963186898915, atol=1e-4), lp)
        model.getCheckpointFiles()

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

        # Test toLocal()
        localModel = model.toLocal()
        self.assertIsInstance(localModel, LDAModel)
        self.assertIsInstance(localModel, LocalLDAModel)
        self.assertNotIsInstance(localModel, DistributedLDAModel)
        self.assertFalse(localModel.isDistributed())
        output = localModel.transform(df)
        expected_cols = ["id", "features", "topicDistribution"]
        self.assertEqual(output.columns, expected_cols)
        self.assertEqual(output.count(), 2)

        # save & load
        with tempfile.TemporaryDirectory(prefix="distributed_lda") as d:
            lda.write().overwrite().save(d)
            lda2 = LDA.load(d)
            self.assertEqual(str(lda), str(lda2))

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

    def test_power_iteration_clustering(self):
        spark = self.spark

        data = [
            (1, 0, 0.5),
            (2, 0, 0.5),
            (2, 1, 0.7),
            (3, 0, 0.5),
            (3, 1, 0.7),
            (3, 2, 0.9),
            (4, 0, 0.5),
            (4, 1, 0.7),
            (4, 2, 0.9),
            (4, 3, 1.1),
            (5, 0, 0.5),
            (5, 1, 0.7),
            (5, 2, 0.9),
            (5, 3, 1.1),
            (5, 4, 1.3),
        ]
        df = spark.createDataFrame(data, ["src", "dst", "weight"]).repartition(1)

        pic = PowerIterationClustering(k=2, weightCol="weight")
        pic.setMaxIter(40)

        self.assertEqual(pic.getK(), 2)
        self.assertEqual(pic.getMaxIter(), 40)
        self.assertEqual(pic.getWeightCol(), "weight")

        assignments = pic.assignClusters(df)
        self.assertEqual(assignments.columns, ["id", "cluster"])
        self.assertEqual(assignments.count(), 6)

        # save & load
        with tempfile.TemporaryDirectory(prefix="power_iteration_clustering") as d:
            pic.write().overwrite().save(d)
            pic2 = PowerIterationClustering.load(d)
            self.assertEqual(str(pic), str(pic2))
            self.assertEqual(pic.uid, pic2.uid)
            self.assertEqual(pic.getK(), pic2.getK())
            self.assertEqual(pic.getMaxIter(), pic2.getMaxIter())
            self.assertEqual(pic.getWeightCol(), pic2.getWeightCol())


class ClusteringTests(ClusteringTestsMixin, ReusedSQLTestCase):
    pass


if __name__ == "__main__":
    from pyspark.ml.tests.test_clustering 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)