clustering.py

#
# 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 sys
import warnings

from typing import Any, Dict, List, Optional, TYPE_CHECKING

import numpy as np

from pyspark import since, keyword_only
from pyspark.ml.param.shared import (
    HasMaxIter,
    HasFeaturesCol,
    HasSeed,
    HasPredictionCol,
    HasAggregationDepth,
    HasWeightCol,
    HasTol,
    HasProbabilityCol,
    HasDistanceMeasure,
    HasCheckpointInterval,
    HasSolver,
    HasMaxBlockSizeInMB,
    Param,
    Params,
    TypeConverters,
)
from pyspark.ml.util import (
    JavaMLWritable,
    JavaMLReadable,
    GeneralJavaMLWritable,
    HasTrainingSummary,
    SparkContext,
)
from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams, JavaWrapper
from pyspark.ml.common import inherit_doc, _java2py
from pyspark.ml.stat import MultivariateGaussian
from pyspark.sql import DataFrame
from pyspark.ml.linalg import Vector, Matrix

if TYPE_CHECKING:
    from pyspark.ml._typing import M
    from py4j.java_gateway import JavaObject


__all__ = [
    "BisectingKMeans",
    "BisectingKMeansModel",
    "BisectingKMeansSummary",
    "KMeans",
    "KMeansModel",
    "KMeansSummary",
    "GaussianMixture",
    "GaussianMixtureModel",
    "GaussianMixtureSummary",
    "LDA",
    "LDAModel",
    "LocalLDAModel",
    "DistributedLDAModel",
    "PowerIterationClustering",
]


class ClusteringSummary(JavaWrapper):
    """
    Clustering results for a given model.

    .. versionadded:: 2.1.0
    """

    @property  # type: ignore[misc]
    @since("2.1.0")
    def predictionCol(self) -> str:
        """
        Name for column of predicted clusters in `predictions`.
        """
        return self._call_java("predictionCol")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def predictions(self) -> DataFrame:
        """
        DataFrame produced by the model's `transform` method.
        """
        return self._call_java("predictions")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def featuresCol(self) -> str:
        """
        Name for column of features in `predictions`.
        """
        return self._call_java("featuresCol")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def k(self) -> int:
        """
        The number of clusters the model was trained with.
        """
        return self._call_java("k")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def cluster(self) -> DataFrame:
        """
        DataFrame of predicted cluster centers for each training data point.
        """
        return self._call_java("cluster")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def clusterSizes(self) -> List[int]:
        """
        Size of (number of data points in) each cluster.
        """
        return self._call_java("clusterSizes")

    @property  # type: ignore[misc]
    @since("2.4.0")
    def numIter(self) -> int:
        """
        Number of iterations.
        """
        return self._call_java("numIter")


@inherit_doc
class _GaussianMixtureParams(
    HasMaxIter,
    HasFeaturesCol,
    HasSeed,
    HasPredictionCol,
    HasProbabilityCol,
    HasTol,
    HasAggregationDepth,
    HasWeightCol,
):
    """
    Params for :py:class:`GaussianMixture` and :py:class:`GaussianMixtureModel`.

    .. versionadded:: 3.0.0
    """

    k: Param[int] = Param(
        Params._dummy(),
        "k",
        "Number of independent Gaussians in the mixture model. " + "Must be > 1.",
        typeConverter=TypeConverters.toInt,
    )

    def __init__(self, *args: Any):
        super(_GaussianMixtureParams, self).__init__(*args)
        self._setDefault(k=2, tol=0.01, maxIter=100, aggregationDepth=2)

    @since("2.0.0")
    def getK(self) -> int:
        """
        Gets the value of `k`
        """
        return self.getOrDefault(self.k)


class GaussianMixtureModel(
    JavaModel,
    _GaussianMixtureParams,
    JavaMLWritable,
    JavaMLReadable["GaussianMixtureModel"],
    HasTrainingSummary["GaussianMixtureSummary"],
):
    """
    Model fitted by GaussianMixture.

    .. versionadded:: 2.0.0
    """

    @since("3.0.0")
    def setFeaturesCol(self, value: str) -> "GaussianMixtureModel":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("3.0.0")
    def setPredictionCol(self, value: str) -> "GaussianMixtureModel":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("3.0.0")
    def setProbabilityCol(self, value: str) -> "GaussianMixtureModel":
        """
        Sets the value of :py:attr:`probabilityCol`.
        """
        return self._set(probabilityCol=value)

    @property  # type: ignore[misc]
    @since("2.0.0")
    def weights(self) -> List[float]:
        """
        Weight for each Gaussian distribution in the mixture.
        This is a multinomial probability distribution over the k Gaussians,
        where weights[i] is the weight for Gaussian i, and weights sum to 1.
        """
        return self._call_java("weights")

    @property  # type: ignore[misc]
    @since("3.0.0")
    def gaussians(self) -> List[MultivariateGaussian]:
        """
        Array of :py:class:`MultivariateGaussian` where gaussians[i] represents
        the Multivariate Gaussian (Normal) Distribution for Gaussian i
        """
        sc = SparkContext._active_spark_context
        assert sc is not None and self._java_obj is not None

        jgaussians = self._java_obj.gaussians()
        return [
            MultivariateGaussian(_java2py(sc, jgaussian.mean()), _java2py(sc, jgaussian.cov()))
            for jgaussian in jgaussians
        ]

    @property  # type: ignore[misc]
    @since("2.0.0")
    def gaussiansDF(self) -> DataFrame:
        """
        Retrieve Gaussian distributions as a DataFrame.
        Each row represents a Gaussian Distribution.
        The DataFrame has two columns: mean (Vector) and cov (Matrix).
        """
        return self._call_java("gaussiansDF")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def summary(self) -> "GaussianMixtureSummary":
        """
        Gets summary (cluster assignments, cluster sizes) of the model trained on the
        training set. An exception is thrown if no summary exists.
        """
        if self.hasSummary:
            return GaussianMixtureSummary(super(GaussianMixtureModel, self).summary)
        else:
            raise RuntimeError(
                "No training summary available for this %s" % self.__class__.__name__
            )

    @since("3.0.0")
    def predict(self, value: Vector) -> int:
        """
        Predict label for the given features.
        """
        return self._call_java("predict", value)

    @since("3.0.0")
    def predictProbability(self, value: Vector) -> Vector:
        """
        Predict probability for the given features.
        """
        return self._call_java("predictProbability", value)


@inherit_doc
class GaussianMixture(
    JavaEstimator[GaussianMixtureModel],
    _GaussianMixtureParams,
    JavaMLWritable,
    JavaMLReadable["GaussianMixture"],
):
    """
    GaussianMixture clustering.
    This class performs expectation maximization for multivariate Gaussian
    Mixture Models (GMMs).  A GMM represents a composite distribution of
    independent Gaussian distributions with associated "mixing" weights
    specifying each's contribution to the composite.

    Given a set of sample points, this class will maximize the log-likelihood
    for a mixture of k Gaussians, iterating until the log-likelihood changes by
    less than convergenceTol, or until it has reached the max number of iterations.
    While this process is generally guaranteed to converge, it is not guaranteed
    to find a global optimum.

    .. versionadded:: 2.0.0

    Notes
    -----
    For high-dimensional data (with many features), this algorithm may perform poorly.
    This is due to high-dimensional data (a) making it difficult to cluster at all
    (based on statistical/theoretical arguments) and (b) numerical issues with
    Gaussian distributions.

    Examples
    --------
    >>> from pyspark.ml.linalg import Vectors

    >>> data = [(Vectors.dense([-0.1, -0.05 ]),),
    ...         (Vectors.dense([-0.01, -0.1]),),
    ...         (Vectors.dense([0.9, 0.8]),),
    ...         (Vectors.dense([0.75, 0.935]),),
    ...         (Vectors.dense([-0.83, -0.68]),),
    ...         (Vectors.dense([-0.91, -0.76]),)]
    >>> df = spark.createDataFrame(data, ["features"])
    >>> gm = GaussianMixture(k=3, tol=0.0001, seed=10)
    >>> gm.getMaxIter()
    100
    >>> gm.setMaxIter(30)
    GaussianMixture...
    >>> gm.getMaxIter()
    30
    >>> model = gm.fit(df)
    >>> model.getAggregationDepth()
    2
    >>> model.getFeaturesCol()
    'features'
    >>> model.setPredictionCol("newPrediction")
    GaussianMixtureModel...
    >>> model.predict(df.head().features)
    2
    >>> model.predictProbability(df.head().features)
    DenseVector([0.0, 0.0, 1.0])
    >>> model.hasSummary
    True
    >>> summary = model.summary
    >>> summary.k
    3
    >>> summary.clusterSizes
    [2, 2, 2]
    >>> weights = model.weights
    >>> len(weights)
    3
    >>> gaussians = model.gaussians
    >>> len(gaussians)
    3
    >>> gaussians[0].mean
    DenseVector([0.825, 0.8675])
    >>> gaussians[0].cov
    DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], 0)
    >>> gaussians[1].mean
    DenseVector([-0.87, -0.72])
    >>> gaussians[1].cov
    DenseMatrix(2, 2, [0.0016, 0.0016, 0.0016, 0.0016], 0)
    >>> gaussians[2].mean
    DenseVector([-0.055, -0.075])
    >>> gaussians[2].cov
    DenseMatrix(2, 2, [0.002, -0.0011, -0.0011, 0.0006], 0)
    >>> model.gaussiansDF.select("mean").head()
    Row(mean=DenseVector([0.825, 0.8675]))
    >>> model.gaussiansDF.select("cov").head()
    Row(cov=DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], False))
    >>> transformed = model.transform(df).select("features", "newPrediction")
    >>> rows = transformed.collect()
    >>> rows[4].newPrediction == rows[5].newPrediction
    True
    >>> rows[2].newPrediction == rows[3].newPrediction
    True
    >>> gmm_path = temp_path + "/gmm"
    >>> gm.save(gmm_path)
    >>> gm2 = GaussianMixture.load(gmm_path)
    >>> gm2.getK()
    3
    >>> model_path = temp_path + "/gmm_model"
    >>> model.save(model_path)
    >>> model2 = GaussianMixtureModel.load(model_path)
    >>> model2.hasSummary
    False
    >>> model2.weights == model.weights
    True
    >>> model2.gaussians[0].mean == model.gaussians[0].mean
    True
    >>> model2.gaussians[0].cov == model.gaussians[0].cov
    True
    >>> model2.gaussians[1].mean == model.gaussians[1].mean
    True
    >>> model2.gaussians[1].cov == model.gaussians[1].cov
    True
    >>> model2.gaussians[2].mean == model.gaussians[2].mean
    True
    >>> model2.gaussians[2].cov == model.gaussians[2].cov
    True
    >>> model2.gaussiansDF.select("mean").head()
    Row(mean=DenseVector([0.825, 0.8675]))
    >>> model2.gaussiansDF.select("cov").head()
    Row(cov=DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], False))
    >>> model.transform(df).take(1) == model2.transform(df).take(1)
    True
    >>> gm2.setWeightCol("weight")
    GaussianMixture...
    """

    _input_kwargs: Dict[str, Any]

    @keyword_only
    def __init__(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        k: int = 2,
        probabilityCol: str = "probability",
        tol: float = 0.01,
        maxIter: int = 100,
        seed: Optional[int] = None,
        aggregationDepth: int = 2,
        weightCol: Optional[str] = None,
    ):
        """
        __init__(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \
                 probabilityCol="probability", tol=0.01, maxIter=100, seed=None, \
                 aggregationDepth=2, weightCol=None)
        """
        super(GaussianMixture, self).__init__()
        self._java_obj = self._new_java_obj(
            "org.apache.spark.ml.clustering.GaussianMixture", self.uid
        )
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def _create_model(self, java_model: "JavaObject") -> "GaussianMixtureModel":
        return GaussianMixtureModel(java_model)

    @keyword_only
    @since("2.0.0")
    def setParams(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        k: int = 2,
        probabilityCol: str = "probability",
        tol: float = 0.01,
        maxIter: int = 100,
        seed: Optional[int] = None,
        aggregationDepth: int = 2,
        weightCol: Optional[str] = None,
    ) -> "GaussianMixture":
        """
        setParams(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \
                  probabilityCol="probability", tol=0.01, maxIter=100, seed=None, \
                  aggregationDepth=2, weightCol=None)

        Sets params for GaussianMixture.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    @since("2.0.0")
    def setK(self, value: int) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`k`.
        """
        return self._set(k=value)

    @since("2.0.0")
    def setMaxIter(self, value: int) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`maxIter`.
        """
        return self._set(maxIter=value)

    @since("2.0.0")
    def setFeaturesCol(self, value: str) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("2.0.0")
    def setPredictionCol(self, value: str) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("2.0.0")
    def setProbabilityCol(self, value: str) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`probabilityCol`.
        """
        return self._set(probabilityCol=value)

    @since("3.0.0")
    def setWeightCol(self, value: str) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`weightCol`.
        """
        return self._set(weightCol=value)

    @since("2.0.0")
    def setSeed(self, value: int) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`seed`.
        """
        return self._set(seed=value)

    @since("2.0.0")
    def setTol(self, value: float) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`tol`.
        """
        return self._set(tol=value)

    @since("3.0.0")
    def setAggregationDepth(self, value: int) -> "GaussianMixture":
        """
        Sets the value of :py:attr:`aggregationDepth`.
        """
        return self._set(aggregationDepth=value)


class GaussianMixtureSummary(ClusteringSummary):
    """
    Gaussian mixture clustering results for a given model.

    .. versionadded:: 2.1.0
    """

    @property  # type: ignore[misc]
    @since("2.1.0")
    def probabilityCol(self) -> str:
        """
        Name for column of predicted probability of each cluster in `predictions`.
        """
        return self._call_java("probabilityCol")

    @property  # type: ignore[misc]
    @since("2.1.0")
    def probability(self) -> DataFrame:
        """
        DataFrame of probabilities of each cluster for each training data point.
        """
        return self._call_java("probability")

    @property  # type: ignore[misc]
    @since("2.2.0")
    def logLikelihood(self) -> float:
        """
        Total log-likelihood for this model on the given data.
        """
        return self._call_java("logLikelihood")


class KMeansSummary(ClusteringSummary):
    """
    Summary of KMeans.

    .. versionadded:: 2.1.0
    """

    @property  # type: ignore[misc]
    @since("2.4.0")
    def trainingCost(self) -> float:
        """
        K-means cost (sum of squared distances to the nearest centroid for all points in the
        training dataset). This is equivalent to sklearn's inertia.
        """
        return self._call_java("trainingCost")


@inherit_doc
class _KMeansParams(
    HasMaxIter,
    HasFeaturesCol,
    HasSeed,
    HasPredictionCol,
    HasTol,
    HasDistanceMeasure,
    HasWeightCol,
    HasSolver,
    HasMaxBlockSizeInMB,
):
    """
    Params for :py:class:`KMeans` and :py:class:`KMeansModel`.

    .. versionadded:: 3.0.0
    """

    k: Param[int] = Param(
        Params._dummy(),
        "k",
        "The number of clusters to create. Must be > 1.",
        typeConverter=TypeConverters.toInt,
    )
    initMode: Param[str] = Param(
        Params._dummy(),
        "initMode",
        'The initialization algorithm. This can be either "random" to '
        + 'choose random points as initial cluster centers, or "k-means||" '
        + "to use a parallel variant of k-means++",
        typeConverter=TypeConverters.toString,
    )
    initSteps: Param[int] = Param(
        Params._dummy(),
        "initSteps",
        "The number of steps for k-means|| " + "initialization mode. Must be > 0.",
        typeConverter=TypeConverters.toInt,
    )
    solver: Param[str] = Param(
        Params._dummy(),
        "solver",
        "The solver algorithm for optimization. Supported " + "options: auto, row, block.",
        typeConverter=TypeConverters.toString,
    )

    def __init__(self, *args: Any):
        super(_KMeansParams, self).__init__(*args)
        self._setDefault(
            k=2,
            initMode="k-means||",
            initSteps=2,
            tol=1e-4,
            maxIter=20,
            distanceMeasure="euclidean",
            solver="auto",
            maxBlockSizeInMB=0.0,
        )

    @since("1.5.0")
    def getK(self) -> int:
        """
        Gets the value of `k`
        """
        return self.getOrDefault(self.k)

    @since("1.5.0")
    def getInitMode(self) -> str:
        """
        Gets the value of `initMode`
        """
        return self.getOrDefault(self.initMode)

    @since("1.5.0")
    def getInitSteps(self) -> int:
        """
        Gets the value of `initSteps`
        """
        return self.getOrDefault(self.initSteps)


class KMeansModel(
    JavaModel,
    _KMeansParams,
    GeneralJavaMLWritable,
    JavaMLReadable["KMeansModel"],
    HasTrainingSummary["KMeansSummary"],
):
    """
    Model fitted by KMeans.

    .. versionadded:: 1.5.0
    """

    @since("3.0.0")
    def setFeaturesCol(self, value: str) -> "KMeansModel":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("3.0.0")
    def setPredictionCol(self, value: str) -> "KMeansModel":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("1.5.0")
    def clusterCenters(self) -> List[np.ndarray]:
        """Get the cluster centers, represented as a list of NumPy arrays."""
        return [c.toArray() for c in self._call_java("clusterCenters")]

    @property  # type: ignore[misc]
    @since("2.1.0")
    def summary(self) -> KMeansSummary:
        """
        Gets summary (cluster assignments, cluster sizes) of the model trained on the
        training set. An exception is thrown if no summary exists.
        """
        if self.hasSummary:
            return KMeansSummary(super(KMeansModel, self).summary)
        else:
            raise RuntimeError(
                "No training summary available for this %s" % self.__class__.__name__
            )

    @since("3.0.0")
    def predict(self, value: Vector) -> int:
        """
        Predict label for the given features.
        """
        return self._call_java("predict", value)


@inherit_doc
class KMeans(JavaEstimator[KMeansModel], _KMeansParams, JavaMLWritable, JavaMLReadable["KMeans"]):
    """
    K-means clustering with a k-means++ like initialization mode
    (the k-means|| algorithm by Bahmani et al).

    .. versionadded:: 1.5.0

    Examples
    --------
    >>> from pyspark.ml.linalg import Vectors
    >>> data = [(Vectors.dense([0.0, 0.0]), 2.0), (Vectors.dense([1.0, 1.0]), 2.0),
    ...         (Vectors.dense([9.0, 8.0]), 2.0), (Vectors.dense([8.0, 9.0]), 2.0)]
    >>> df = spark.createDataFrame(data, ["features", "weighCol"])
    >>> kmeans = KMeans(k=2)
    >>> kmeans.setSeed(1)
    KMeans...
    >>> kmeans.setWeightCol("weighCol")
    KMeans...
    >>> kmeans.setMaxIter(10)
    KMeans...
    >>> kmeans.getMaxIter()
    10
    >>> kmeans.clear(kmeans.maxIter)
    >>> kmeans.getSolver()
    'auto'
    >>> model = kmeans.fit(df)
    >>> model.getMaxBlockSizeInMB()
    0.0
    >>> model.getDistanceMeasure()
    'euclidean'
    >>> model.setPredictionCol("newPrediction")
    KMeansModel...
    >>> model.predict(df.head().features)
    0
    >>> centers = model.clusterCenters()
    >>> len(centers)
    2
    >>> transformed = model.transform(df).select("features", "newPrediction")
    >>> rows = transformed.collect()
    >>> rows[0].newPrediction == rows[1].newPrediction
    True
    >>> rows[2].newPrediction == rows[3].newPrediction
    True
    >>> model.hasSummary
    True
    >>> summary = model.summary
    >>> summary.k
    2
    >>> summary.clusterSizes
    [2, 2]
    >>> summary.trainingCost
    4.0
    >>> kmeans_path = temp_path + "/kmeans"
    >>> kmeans.save(kmeans_path)
    >>> kmeans2 = KMeans.load(kmeans_path)
    >>> kmeans2.getK()
    2
    >>> model_path = temp_path + "/kmeans_model"
    >>> model.save(model_path)
    >>> model2 = KMeansModel.load(model_path)
    >>> model2.hasSummary
    False
    >>> model.clusterCenters()[0] == model2.clusterCenters()[0]
    array([ True,  True], dtype=bool)
    >>> model.clusterCenters()[1] == model2.clusterCenters()[1]
    array([ True,  True], dtype=bool)
    >>> model.transform(df).take(1) == model2.transform(df).take(1)
    True
    """

    _input_kwargs: Dict[str, Any]

    @keyword_only
    def __init__(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        k: int = 2,
        initMode: str = "k-means||",
        initSteps: int = 2,
        tol: float = 1e-4,
        maxIter: int = 20,
        seed: Optional[int] = None,
        distanceMeasure: str = "euclidean",
        weightCol: Optional[str] = None,
        solver: str = "auto",
        maxBlockSizeInMB: float = 0.0,
    ):
        """
        __init__(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \
                 initMode="k-means||", initSteps=2, tol=1e-4, maxIter=20, seed=None, \
                 distanceMeasure="euclidean", weightCol=None, solver="auto", \
                 maxBlockSizeInMB=0.0)
        """
        super(KMeans, self).__init__()
        self._java_obj = self._new_java_obj("org.apache.spark.ml.clustering.KMeans", self.uid)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def _create_model(self, java_model: "JavaObject") -> KMeansModel:
        return KMeansModel(java_model)

    @keyword_only
    @since("1.5.0")
    def setParams(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        k: int = 2,
        initMode: str = "k-means||",
        initSteps: int = 2,
        tol: float = 1e-4,
        maxIter: int = 20,
        seed: Optional[int] = None,
        distanceMeasure: str = "euclidean",
        weightCol: Optional[str] = None,
        solver: str = "auto",
        maxBlockSizeInMB: float = 0.0,
    ) -> "KMeans":
        """
        setParams(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \
                  initMode="k-means||", initSteps=2, tol=1e-4, maxIter=20, seed=None, \
                  distanceMeasure="euclidean", weightCol=None, solver="auto", \
                  maxBlockSizeInMB=0.0)

        Sets params for KMeans.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    @since("1.5.0")
    def setK(self, value: int) -> "KMeans":
        """
        Sets the value of :py:attr:`k`.
        """
        return self._set(k=value)

    @since("1.5.0")
    def setInitMode(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`initMode`.
        """
        return self._set(initMode=value)

    @since("1.5.0")
    def setInitSteps(self, value: int) -> "KMeans":
        """
        Sets the value of :py:attr:`initSteps`.
        """
        return self._set(initSteps=value)

    @since("2.4.0")
    def setDistanceMeasure(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`distanceMeasure`.
        """
        return self._set(distanceMeasure=value)

    @since("1.5.0")
    def setMaxIter(self, value: int) -> "KMeans":
        """
        Sets the value of :py:attr:`maxIter`.
        """
        return self._set(maxIter=value)

    @since("1.5.0")
    def setFeaturesCol(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("1.5.0")
    def setPredictionCol(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("1.5.0")
    def setSeed(self, value: int) -> "KMeans":
        """
        Sets the value of :py:attr:`seed`.
        """
        return self._set(seed=value)

    @since("1.5.0")
    def setTol(self, value: float) -> "KMeans":
        """
        Sets the value of :py:attr:`tol`.
        """
        return self._set(tol=value)

    @since("3.0.0")
    def setWeightCol(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`weightCol`.
        """
        return self._set(weightCol=value)

    @since("3.4.0")
    def setSolver(self, value: str) -> "KMeans":
        """
        Sets the value of :py:attr:`solver`.
        """
        return self._set(solver=value)

    @since("3.4.0")
    def setMaxBlockSizeInMB(self, value: float) -> "KMeans":
        """
        Sets the value of :py:attr:`maxBlockSizeInMB`.
        """
        return self._set(maxBlockSizeInMB=value)


@inherit_doc
class _BisectingKMeansParams(
    HasMaxIter,
    HasFeaturesCol,
    HasSeed,
    HasPredictionCol,
    HasDistanceMeasure,
    HasWeightCol,
):
    """
    Params for :py:class:`BisectingKMeans` and :py:class:`BisectingKMeansModel`.

    .. versionadded:: 3.0.0
    """

    k: Param[int] = Param(
        Params._dummy(),
        "k",
        "The desired number of leaf clusters. Must be > 1.",
        typeConverter=TypeConverters.toInt,
    )
    minDivisibleClusterSize: Param[float] = Param(
        Params._dummy(),
        "minDivisibleClusterSize",
        "The minimum number of points (if >= 1.0) or the minimum "
        + "proportion of points (if < 1.0) of a divisible cluster.",
        typeConverter=TypeConverters.toFloat,
    )

    def __init__(self, *args: Any):
        super(_BisectingKMeansParams, self).__init__(*args)
        self._setDefault(maxIter=20, k=4, minDivisibleClusterSize=1.0)

    @since("2.0.0")
    def getK(self) -> int:
        """
        Gets the value of `k` or its default value.
        """
        return self.getOrDefault(self.k)

    @since("2.0.0")
    def getMinDivisibleClusterSize(self) -> float:
        """
        Gets the value of `minDivisibleClusterSize` or its default value.
        """
        return self.getOrDefault(self.minDivisibleClusterSize)


class BisectingKMeansModel(
    JavaModel,
    _BisectingKMeansParams,
    JavaMLWritable,
    JavaMLReadable["BisectingKMeansModel"],
    HasTrainingSummary["BisectingKMeansSummary"],
):
    """
    Model fitted by BisectingKMeans.

    .. versionadded:: 2.0.0
    """

    @since("3.0.0")
    def setFeaturesCol(self, value: str) -> "BisectingKMeansModel":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("3.0.0")
    def setPredictionCol(self, value: str) -> "BisectingKMeansModel":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("2.0.0")
    def clusterCenters(self) -> List[np.ndarray]:
        """Get the cluster centers, represented as a list of NumPy arrays."""
        return [c.toArray() for c in self._call_java("clusterCenters")]

    @since("2.0.0")
    def computeCost(self, dataset: DataFrame) -> float:
        """
        Computes the sum of squared distances between the input points
        and their corresponding cluster centers.

        .. deprecated:: 3.0.0
            It will be removed in future versions. Use :py:class:`ClusteringEvaluator` instead.
            You can also get the cost on the training dataset in the summary.
        """
        warnings.warn(
            "Deprecated in 3.0.0. It will be removed in future versions. Use "
            "ClusteringEvaluator instead. You can also get the cost on the training "
            "dataset in the summary.",
            FutureWarning,
        )
        return self._call_java("computeCost", dataset)

    @property  # type: ignore[misc]
    @since("2.1.0")
    def summary(self) -> "BisectingKMeansSummary":
        """
        Gets summary (cluster assignments, cluster sizes) of the model trained on the
        training set. An exception is thrown if no summary exists.
        """
        if self.hasSummary:
            return BisectingKMeansSummary(super(BisectingKMeansModel, self).summary)
        else:
            raise RuntimeError(
                "No training summary available for this %s" % self.__class__.__name__
            )

    @since("3.0.0")
    def predict(self, value: Vector) -> int:
        """
        Predict label for the given features.
        """
        return self._call_java("predict", value)


@inherit_doc
class BisectingKMeans(
    JavaEstimator[BisectingKMeansModel],
    _BisectingKMeansParams,
    JavaMLWritable,
    JavaMLReadable["BisectingKMeans"],
):
    """
    A bisecting k-means algorithm based on the paper "A comparison of document clustering
    techniques" by Steinbach, Karypis, and Kumar, with modification to fit Spark.
    The algorithm starts from a single cluster that contains all points.
    Iteratively it finds divisible clusters on the bottom level and bisects each of them using
    k-means, until there are `k` leaf clusters in total or no leaf clusters are divisible.
    The bisecting steps of clusters on the same level are grouped together to increase parallelism.
    If bisecting all divisible clusters on the bottom level would result more than `k` leaf
    clusters, larger clusters get higher priority.

    .. versionadded:: 2.0.0

    Examples
    --------
    >>> from pyspark.ml.linalg import Vectors
    >>> data = [(Vectors.dense([0.0, 0.0]), 2.0), (Vectors.dense([1.0, 1.0]), 2.0),
    ...         (Vectors.dense([9.0, 8.0]), 2.0), (Vectors.dense([8.0, 9.0]), 2.0)]
    >>> df = spark.createDataFrame(data, ["features", "weighCol"])
    >>> bkm = BisectingKMeans(k=2, minDivisibleClusterSize=1.0)
    >>> bkm.setMaxIter(10)
    BisectingKMeans...
    >>> bkm.getMaxIter()
    10
    >>> bkm.clear(bkm.maxIter)
    >>> bkm.setSeed(1)
    BisectingKMeans...
    >>> bkm.setWeightCol("weighCol")
    BisectingKMeans...
    >>> bkm.getSeed()
    1
    >>> bkm.clear(bkm.seed)
    >>> model = bkm.fit(df)
    >>> model.getMaxIter()
    20
    >>> model.setPredictionCol("newPrediction")
    BisectingKMeansModel...
    >>> model.predict(df.head().features)
    0
    >>> centers = model.clusterCenters()
    >>> len(centers)
    2
    >>> model.computeCost(df)
    2.0
    >>> model.hasSummary
    True
    >>> summary = model.summary
    >>> summary.k
    2
    >>> summary.clusterSizes
    [2, 2]
    >>> summary.trainingCost
    4.000...
    >>> transformed = model.transform(df).select("features", "newPrediction")
    >>> rows = transformed.collect()
    >>> rows[0].newPrediction == rows[1].newPrediction
    True
    >>> rows[2].newPrediction == rows[3].newPrediction
    True
    >>> bkm_path = temp_path + "/bkm"
    >>> bkm.save(bkm_path)
    >>> bkm2 = BisectingKMeans.load(bkm_path)
    >>> bkm2.getK()
    2
    >>> bkm2.getDistanceMeasure()
    'euclidean'
    >>> model_path = temp_path + "/bkm_model"
    >>> model.save(model_path)
    >>> model2 = BisectingKMeansModel.load(model_path)
    >>> model2.hasSummary
    False
    >>> model.clusterCenters()[0] == model2.clusterCenters()[0]
    array([ True,  True], dtype=bool)
    >>> model.clusterCenters()[1] == model2.clusterCenters()[1]
    array([ True,  True], dtype=bool)
    >>> model.transform(df).take(1) == model2.transform(df).take(1)
    True
    """

    _input_kwargs: Dict[str, Any]

    @keyword_only
    def __init__(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        maxIter: int = 20,
        seed: Optional[int] = None,
        k: int = 4,
        minDivisibleClusterSize: float = 1.0,
        distanceMeasure: str = "euclidean",
        weightCol: Optional[str] = None,
    ):
        """
        __init__(self, \\*, featuresCol="features", predictionCol="prediction", maxIter=20, \
                 seed=None, k=4, minDivisibleClusterSize=1.0, distanceMeasure="euclidean", \
                 weightCol=None)
        """
        super(BisectingKMeans, self).__init__()
        self._java_obj = self._new_java_obj(
            "org.apache.spark.ml.clustering.BisectingKMeans", self.uid
        )
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    @since("2.0.0")
    def setParams(
        self,
        *,
        featuresCol: str = "features",
        predictionCol: str = "prediction",
        maxIter: int = 20,
        seed: Optional[int] = None,
        k: int = 4,
        minDivisibleClusterSize: float = 1.0,
        distanceMeasure: str = "euclidean",
        weightCol: Optional[str] = None,
    ) -> "BisectingKMeans":
        """
        setParams(self, \\*, featuresCol="features", predictionCol="prediction", maxIter=20, \
                  seed=None, k=4, minDivisibleClusterSize=1.0, distanceMeasure="euclidean", \
                  weightCol=None)
        Sets params for BisectingKMeans.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    @since("2.0.0")
    def setK(self, value: int) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`k`.
        """
        return self._set(k=value)

    @since("2.0.0")
    def setMinDivisibleClusterSize(self, value: float) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`minDivisibleClusterSize`.
        """
        return self._set(minDivisibleClusterSize=value)

    @since("2.4.0")
    def setDistanceMeasure(self, value: str) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`distanceMeasure`.
        """
        return self._set(distanceMeasure=value)

    @since("2.0.0")
    def setMaxIter(self, value: int) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`maxIter`.
        """
        return self._set(maxIter=value)

    @since("2.0.0")
    def setFeaturesCol(self, value: str) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("2.0.0")
    def setPredictionCol(self, value: str) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`predictionCol`.
        """
        return self._set(predictionCol=value)

    @since("2.0.0")
    def setSeed(self, value: int) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`seed`.
        """
        return self._set(seed=value)

    @since("3.0.0")
    def setWeightCol(self, value: str) -> "BisectingKMeans":
        """
        Sets the value of :py:attr:`weightCol`.
        """
        return self._set(weightCol=value)

    def _create_model(self, java_model: "JavaObject") -> BisectingKMeansModel:
        return BisectingKMeansModel(java_model)


class BisectingKMeansSummary(ClusteringSummary):
    """
    Bisecting KMeans clustering results for a given model.

    .. versionadded:: 2.1.0
    """

    @property  # type: ignore[misc]
    @since("3.0.0")
    def trainingCost(self) -> float:
        """
        Sum of squared distances to the nearest centroid for all points in the training dataset.
        This is equivalent to sklearn's inertia.
        """
        return self._call_java("trainingCost")


@inherit_doc
class _LDAParams(HasMaxIter, HasFeaturesCol, HasSeed, HasCheckpointInterval):
    """
    Params for :py:class:`LDA` and :py:class:`LDAModel`.

    .. versionadded:: 3.0.0
    """

    k: Param[int] = Param(
        Params._dummy(),
        "k",
        "The number of topics (clusters) to infer. Must be > 1.",
        typeConverter=TypeConverters.toInt,
    )
    optimizer: Param[str] = Param(
        Params._dummy(),
        "optimizer",
        "Optimizer or inference algorithm used to estimate the LDA model.  "
        "Supported: online, em",
        typeConverter=TypeConverters.toString,
    )
    learningOffset: Param[float] = Param(
        Params._dummy(),
        "learningOffset",
        "A (positive) learning parameter that downweights early iterations."
        " Larger values make early iterations count less",
        typeConverter=TypeConverters.toFloat,
    )
    learningDecay: Param[float] = Param(
        Params._dummy(),
        "learningDecay",
        "Learning rate, set as an"
        "exponential decay rate. This should be between (0.5, 1.0] to "
        "guarantee asymptotic convergence.",
        typeConverter=TypeConverters.toFloat,
    )
    subsamplingRate: Param[float] = Param(
        Params._dummy(),
        "subsamplingRate",
        "Fraction of the corpus to be sampled and used in each iteration "
        "of mini-batch gradient descent, in range (0, 1].",
        typeConverter=TypeConverters.toFloat,
    )
    optimizeDocConcentration: Param[bool] = Param(
        Params._dummy(),
        "optimizeDocConcentration",
        "Indicates whether the docConcentration (Dirichlet parameter "
        "for document-topic distribution) will be optimized during "
        "training.",
        typeConverter=TypeConverters.toBoolean,
    )
    docConcentration: Param[List[float]] = Param(
        Params._dummy(),
        "docConcentration",
        'Concentration parameter (commonly named "alpha") for the '
        'prior placed on documents\' distributions over topics ("theta").',
        typeConverter=TypeConverters.toListFloat,
    )
    topicConcentration: Param[float] = Param(
        Params._dummy(),
        "topicConcentration",
        'Concentration parameter (commonly named "beta" or "eta") for '
        "the prior placed on topic' distributions over terms.",
        typeConverter=TypeConverters.toFloat,
    )
    topicDistributionCol: Param[str] = Param(
        Params._dummy(),
        "topicDistributionCol",
        "Output column with estimates of the topic mixture distribution "
        'for each document (often called "theta" in the literature). '
        "Returns a vector of zeros for an empty document.",
        typeConverter=TypeConverters.toString,
    )
    keepLastCheckpoint: Param[bool] = Param(
        Params._dummy(),
        "keepLastCheckpoint",
        "(For EM optimizer) If using checkpointing, this indicates whether"
        " to keep the last checkpoint. If false, then the checkpoint will be"
        " deleted. Deleting the checkpoint can cause failures if a data"
        " partition is lost, so set this bit with care.",
        TypeConverters.toBoolean,
    )

    def __init__(self, *args: Any):
        super(_LDAParams, self).__init__(*args)
        self._setDefault(
            maxIter=20,
            checkpointInterval=10,
            k=10,
            optimizer="online",
            learningOffset=1024.0,
            learningDecay=0.51,
            subsamplingRate=0.05,
            optimizeDocConcentration=True,
            topicDistributionCol="topicDistribution",
            keepLastCheckpoint=True,
        )

    @since("2.0.0")
    def getK(self) -> int:
        """
        Gets the value of :py:attr:`k` or its default value.
        """
        return self.getOrDefault(self.k)

    @since("2.0.0")
    def getOptimizer(self) -> str:
        """
        Gets the value of :py:attr:`optimizer` or its default value.
        """
        return self.getOrDefault(self.optimizer)

    @since("2.0.0")
    def getLearningOffset(self) -> float:
        """
        Gets the value of :py:attr:`learningOffset` or its default value.
        """
        return self.getOrDefault(self.learningOffset)

    @since("2.0.0")
    def getLearningDecay(self) -> float:
        """
        Gets the value of :py:attr:`learningDecay` or its default value.
        """
        return self.getOrDefault(self.learningDecay)

    @since("2.0.0")
    def getSubsamplingRate(self) -> float:
        """
        Gets the value of :py:attr:`subsamplingRate` or its default value.
        """
        return self.getOrDefault(self.subsamplingRate)

    @since("2.0.0")
    def getOptimizeDocConcentration(self) -> bool:
        """
        Gets the value of :py:attr:`optimizeDocConcentration` or its default value.
        """
        return self.getOrDefault(self.optimizeDocConcentration)

    @since("2.0.0")
    def getDocConcentration(self) -> List[float]:
        """
        Gets the value of :py:attr:`docConcentration` or its default value.
        """
        return self.getOrDefault(self.docConcentration)

    @since("2.0.0")
    def getTopicConcentration(self) -> float:
        """
        Gets the value of :py:attr:`topicConcentration` or its default value.
        """
        return self.getOrDefault(self.topicConcentration)

    @since("2.0.0")
    def getTopicDistributionCol(self) -> str:
        """
        Gets the value of :py:attr:`topicDistributionCol` or its default value.
        """
        return self.getOrDefault(self.topicDistributionCol)

    @since("2.0.0")
    def getKeepLastCheckpoint(self) -> bool:
        """
        Gets the value of :py:attr:`keepLastCheckpoint` or its default value.
        """
        return self.getOrDefault(self.keepLastCheckpoint)


@inherit_doc
class LDAModel(JavaModel, _LDAParams):
    """
    Latent Dirichlet Allocation (LDA) model.
    This abstraction permits for different underlying representations,
    including local and distributed data structures.

    .. versionadded:: 2.0.0
    """

    @since("3.0.0")
    def setFeaturesCol(self: "M", value: str) -> "M":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)

    @since("3.0.0")
    def setSeed(self: "M", value: int) -> "M":
        """
        Sets the value of :py:attr:`seed`.
        """
        return self._set(seed=value)

    @since("3.0.0")
    def setTopicDistributionCol(self: "M", value: str) -> "M":
        """
        Sets the value of :py:attr:`topicDistributionCol`.
        """
        return self._set(topicDistributionCol=value)

    @since("2.0.0")
    def isDistributed(self) -> bool:
        """
        Indicates whether this instance is of type DistributedLDAModel
        """
        return self._call_java("isDistributed")

    @since("2.0.0")
    def vocabSize(self) -> int:
        """Vocabulary size (number of terms or words in the vocabulary)"""
        return self._call_java("vocabSize")

    @since("2.0.0")
    def topicsMatrix(self) -> Matrix:
        """
        Inferred topics, where each topic is represented by a distribution over terms.
        This is a matrix of size vocabSize x k, where each column is a topic.
        No guarantees are given about the ordering of the topics.

        .. warning:: If this model is actually a :py:class:`DistributedLDAModel`
            instance produced by the Expectation-Maximization ("em") `optimizer`,
            then this method could involve collecting a large amount of data
            to the driver (on the order of vocabSize x k).
        """
        return self._call_java("topicsMatrix")

    @since("2.0.0")
    def logLikelihood(self, dataset: DataFrame) -> float:
        """
        Calculates a lower bound on the log likelihood of the entire corpus.
        See Equation (16) in the Online LDA paper (Hoffman et al., 2010).

        .. warning:: If this model is an instance of :py:class:`DistributedLDAModel` (produced when
            :py:attr:`optimizer` is set to "em"), this involves collecting a large
            :py:func:`topicsMatrix` to the driver. This implementation may be changed in the future.
        """
        return self._call_java("logLikelihood", dataset)

    @since("2.0.0")
    def logPerplexity(self, dataset: DataFrame) -> float:
        """
        Calculate an upper bound on perplexity.  (Lower is better.)
        See Equation (16) in the Online LDA paper (Hoffman et al., 2010).

        .. warning:: If this model is an instance of :py:class:`DistributedLDAModel` (produced when
            :py:attr:`optimizer` is set to "em"), this involves collecting a large
            :py:func:`topicsMatrix` to the driver. This implementation may be changed in the future.
        """
        return self._call_java("logPerplexity", dataset)

    @since("2.0.0")
    def describeTopics(self, maxTermsPerTopic: int = 10) -> DataFrame:
        """
        Return the topics described by their top-weighted terms.
        """
        return self._call_java("describeTopics", maxTermsPerTopic)

    @since("2.0.0")
    def estimatedDocConcentration(self) -> Vector:
        """
        Value for :py:attr:`LDA.docConcentration` estimated from data.
        If Online LDA was used and :py:attr:`LDA.optimizeDocConcentration` was set to false,
        then this returns the fixed (given) value for the :py:attr:`LDA.docConcentration` parameter.
        """
        return self._call_java("estimatedDocConcentration")


@inherit_doc
class DistributedLDAModel(LDAModel, JavaMLReadable["DistributedLDAModel"], JavaMLWritable):
    """
    Distributed model fitted by :py:class:`LDA`.
    This type of model is currently only produced by Expectation-Maximization (EM).

    This model stores the inferred topics, the full training dataset, and the topic distribution
    for each training document.

    .. versionadded:: 2.0.0
    """

    @since("2.0.0")
    def toLocal(self) -> "LocalLDAModel":
        """
        Convert this distributed model to a local representation.  This discards info about the
        training dataset.

        .. warning:: This involves collecting a large :py:func:`topicsMatrix` to the driver.
        """
        model = LocalLDAModel(self._call_java("toLocal"))

        # SPARK-10931: Temporary fix to be removed once LDAModel defines Params
        model._create_params_from_java()
        model._transfer_params_from_java()

        return model

    @since("2.0.0")
    def trainingLogLikelihood(self) -> float:
        """
        Log likelihood of the observed tokens in the training set,
        given the current parameter estimates:
        log P(docs | topics, topic distributions for docs, Dirichlet hyperparameters)

        Notes
        -----
        - This excludes the prior; for that, use :py:func:`logPrior`.
        - Even with :py:func:`logPrior`, this is NOT the same as the data log likelihood given
            the hyperparameters.
        - This is computed from the topic distributions computed during training. If you call
            :py:func:`logLikelihood` on the same training dataset, the topic distributions
            will be computed again, possibly giving different results.
        """
        return self._call_java("trainingLogLikelihood")

    @since("2.0.0")
    def logPrior(self) -> float:
        """
        Log probability of the current parameter estimate:
        log P(topics, topic distributions for docs | alpha, eta)
        """
        return self._call_java("logPrior")

    def getCheckpointFiles(self) -> List[str]:
        """
        If using checkpointing and :py:attr:`LDA.keepLastCheckpoint` is set to true, then there may
        be saved checkpoint files.  This method is provided so that users can manage those files.

        .. versionadded:: 2.0.0

        Returns
        -------
        list
            List of checkpoint files from training

        Notes
        -----
        Removing the checkpoints can cause failures if a partition is lost and is needed
        by certain :py:class:`DistributedLDAModel` methods.  Reference counting will clean up
        the checkpoints when this model and derivative data go out of scope.
        """
        return self._call_java("getCheckpointFiles")


@inherit_doc
class LocalLDAModel(LDAModel, JavaMLReadable["LocalLDAModel"], JavaMLWritable):
    """
    Local (non-distributed) model fitted by :py:class:`LDA`.
    This model stores the inferred topics only; it does not store info about the training dataset.

    .. versionadded:: 2.0.0
    """

    pass


@inherit_doc
class LDA(JavaEstimator[LDAModel], _LDAParams, JavaMLReadable["LDA"], JavaMLWritable):
    """
    Latent Dirichlet Allocation (LDA), a topic model designed for text documents.

    Terminology:

     - "term" = "word": an element of the vocabulary
     - "token": instance of a term appearing in a document
     - "topic": multinomial distribution over terms representing some concept
     - "document": one piece of text, corresponding to one row in the input data

    Original LDA paper (journal version):
      Blei, Ng, and Jordan.  "Latent Dirichlet Allocation."  JMLR, 2003.

    Input data (featuresCol):
    LDA is given a collection of documents as input data, via the featuresCol parameter.
    Each document is specified as a :py:class:`Vector` of length vocabSize, where each entry is the
    count for the corresponding term (word) in the document.  Feature transformers such as
    :py:class:`pyspark.ml.feature.Tokenizer` and :py:class:`pyspark.ml.feature.CountVectorizer`
    can be useful for converting text to word count vectors.

    .. versionadded:: 2.0.0

    Examples
    --------
    >>> from pyspark.ml.linalg import Vectors, SparseVector
    >>> from pyspark.ml.clustering import LDA
    >>> df = spark.createDataFrame([[1, Vectors.dense([0.0, 1.0])],
    ...      [2, SparseVector(2, {0: 1.0})],], ["id", "features"])
    >>> lda = LDA(k=2, seed=1, optimizer="em")
    >>> lda.setMaxIter(10)
    LDA...
    >>> lda.getMaxIter()
    10
    >>> lda.clear(lda.maxIter)
    >>> model = lda.fit(df)
    >>> model.setSeed(1)
    DistributedLDAModel...
    >>> model.getTopicDistributionCol()
    'topicDistribution'
    >>> model.isDistributed()
    True
    >>> localModel = model.toLocal()
    >>> localModel.isDistributed()
    False
    >>> model.vocabSize()
    2
    >>> model.describeTopics().show()
    +-----+-----------+--------------------+
    |topic|termIndices|         termWeights|
    +-----+-----------+--------------------+
    |    0|     [1, 0]|[0.50401530077160...|
    |    1|     [0, 1]|[0.50401530077160...|
    +-----+-----------+--------------------+
    ...
    >>> model.topicsMatrix()
    DenseMatrix(2, 2, [0.496, 0.504, 0.504, 0.496], 0)
    >>> lda_path = temp_path + "/lda"
    >>> lda.save(lda_path)
    >>> sameLDA = LDA.load(lda_path)
    >>> distributed_model_path = temp_path + "/lda_distributed_model"
    >>> model.save(distributed_model_path)
    >>> sameModel = DistributedLDAModel.load(distributed_model_path)
    >>> local_model_path = temp_path + "/lda_local_model"
    >>> localModel.save(local_model_path)
    >>> sameLocalModel = LocalLDAModel.load(local_model_path)
    >>> model.transform(df).take(1) == sameLocalModel.transform(df).take(1)
    True
    """

    _input_kwargs: Dict[str, Any]

    @keyword_only
    def __init__(
        self,
        *,
        featuresCol: str = "features",
        maxIter: int = 20,
        seed: Optional[int] = None,
        checkpointInterval: int = 10,
        k: int = 10,
        optimizer: str = "online",
        learningOffset: float = 1024.0,
        learningDecay: float = 0.51,
        subsamplingRate: float = 0.05,
        optimizeDocConcentration: bool = True,
        docConcentration: Optional[List[float]] = None,
        topicConcentration: Optional[float] = None,
        topicDistributionCol: str = "topicDistribution",
        keepLastCheckpoint: bool = True,
    ):
        """
        __init__(self, \\*, featuresCol="features", maxIter=20, seed=None, checkpointInterval=10,\
                  k=10, optimizer="online", learningOffset=1024.0, learningDecay=0.51,\
                  subsamplingRate=0.05, optimizeDocConcentration=True,\
                  docConcentration=None, topicConcentration=None,\
                  topicDistributionCol="topicDistribution", keepLastCheckpoint=True)
        """
        super(LDA, self).__init__()
        self._java_obj = self._new_java_obj("org.apache.spark.ml.clustering.LDA", self.uid)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def _create_model(self, java_model: "JavaObject") -> LDAModel:
        if self.getOptimizer() == "em":
            return DistributedLDAModel(java_model)
        else:
            return LocalLDAModel(java_model)

    @keyword_only
    @since("2.0.0")
    def setParams(
        self,
        *,
        featuresCol: str = "features",
        maxIter: int = 20,
        seed: Optional[int] = None,
        checkpointInterval: int = 10,
        k: int = 10,
        optimizer: str = "online",
        learningOffset: float = 1024.0,
        learningDecay: float = 0.51,
        subsamplingRate: float = 0.05,
        optimizeDocConcentration: bool = True,
        docConcentration: Optional[List[float]] = None,
        topicConcentration: Optional[float] = None,
        topicDistributionCol: str = "topicDistribution",
        keepLastCheckpoint: bool = True,
    ) -> "LDA":
        """
        setParams(self, \\*, featuresCol="features", maxIter=20, seed=None, checkpointInterval=10,\
                  k=10, optimizer="online", learningOffset=1024.0, learningDecay=0.51,\
                  subsamplingRate=0.05, optimizeDocConcentration=True,\
                  docConcentration=None, topicConcentration=None,\
                  topicDistributionCol="topicDistribution", keepLastCheckpoint=True)

        Sets params for LDA.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    @since("2.0.0")
    def setCheckpointInterval(self, value: int) -> "LDA":
        """
        Sets the value of :py:attr:`checkpointInterval`.
        """
        return self._set(checkpointInterval=value)

    @since("2.0.0")
    def setSeed(self, value: int) -> "LDA":
        """
        Sets the value of :py:attr:`seed`.
        """
        return self._set(seed=value)

    @since("2.0.0")
    def setK(self, value: int) -> "LDA":
        """
        Sets the value of :py:attr:`k`.

        >>> algo = LDA().setK(10)
        >>> algo.getK()
        10
        """
        return self._set(k=value)

    @since("2.0.0")
    def setOptimizer(self, value: str) -> "LDA":
        """
        Sets the value of :py:attr:`optimizer`.
        Currently only support 'em' and 'online'.

        Examples
        --------
        >>> algo = LDA().setOptimizer("em")
        >>> algo.getOptimizer()
        'em'
        """
        return self._set(optimizer=value)

    @since("2.0.0")
    def setLearningOffset(self, value: float) -> "LDA":
        """
        Sets the value of :py:attr:`learningOffset`.

        Examples
        --------
        >>> algo = LDA().setLearningOffset(100)
        >>> algo.getLearningOffset()
        100.0
        """
        return self._set(learningOffset=value)

    @since("2.0.0")
    def setLearningDecay(self, value: float) -> "LDA":
        """
        Sets the value of :py:attr:`learningDecay`.

        Examples
        --------
        >>> algo = LDA().setLearningDecay(0.1)
        >>> algo.getLearningDecay()
        0.1...
        """
        return self._set(learningDecay=value)

    @since("2.0.0")
    def setSubsamplingRate(self, value: float) -> "LDA":
        """
        Sets the value of :py:attr:`subsamplingRate`.

        Examples
        --------
        >>> algo = LDA().setSubsamplingRate(0.1)
        >>> algo.getSubsamplingRate()
        0.1...
        """
        return self._set(subsamplingRate=value)

    @since("2.0.0")
    def setOptimizeDocConcentration(self, value: bool) -> "LDA":
        """
        Sets the value of :py:attr:`optimizeDocConcentration`.

        Examples
        --------
        >>> algo = LDA().setOptimizeDocConcentration(True)
        >>> algo.getOptimizeDocConcentration()
        True
        """
        return self._set(optimizeDocConcentration=value)

    @since("2.0.0")
    def setDocConcentration(self, value: List[float]) -> "LDA":
        """
        Sets the value of :py:attr:`docConcentration`.

        Examples
        --------
        >>> algo = LDA().setDocConcentration([0.1, 0.2])
        >>> algo.getDocConcentration()
        [0.1..., 0.2...]
        """
        return self._set(docConcentration=value)

    @since("2.0.0")
    def setTopicConcentration(self, value: float) -> "LDA":
        """
        Sets the value of :py:attr:`topicConcentration`.

        Examples
        --------
        >>> algo = LDA().setTopicConcentration(0.5)
        >>> algo.getTopicConcentration()
        0.5...
        """
        return self._set(topicConcentration=value)

    @since("2.0.0")
    def setTopicDistributionCol(self, value: str) -> "LDA":
        """
        Sets the value of :py:attr:`topicDistributionCol`.

        Examples
        --------
        >>> algo = LDA().setTopicDistributionCol("topicDistributionCol")
        >>> algo.getTopicDistributionCol()
        'topicDistributionCol'
        """
        return self._set(topicDistributionCol=value)

    @since("2.0.0")
    def setKeepLastCheckpoint(self, value: bool) -> "LDA":
        """
        Sets the value of :py:attr:`keepLastCheckpoint`.

        Examples
        --------
        >>> algo = LDA().setKeepLastCheckpoint(False)
        >>> algo.getKeepLastCheckpoint()
        False
        """
        return self._set(keepLastCheckpoint=value)

    @since("2.0.0")
    def setMaxIter(self, value: int) -> "LDA":
        """
        Sets the value of :py:attr:`maxIter`.
        """
        return self._set(maxIter=value)

    @since("2.0.0")
    def setFeaturesCol(self, value: str) -> "LDA":
        """
        Sets the value of :py:attr:`featuresCol`.
        """
        return self._set(featuresCol=value)


@inherit_doc
class _PowerIterationClusteringParams(HasMaxIter, HasWeightCol):
    """
    Params for :py:class:`PowerIterationClustering`.

    .. versionadded:: 3.0.0
    """

    k: Param[int] = Param(
        Params._dummy(),
        "k",
        "The number of clusters to create. Must be > 1.",
        typeConverter=TypeConverters.toInt,
    )
    initMode: Param[str] = Param(
        Params._dummy(),
        "initMode",
        "The initialization algorithm. This can be either "
        + "'random' to use a random vector as vertex properties, or 'degree' to use "
        + "a normalized sum of similarities with other vertices.  Supported options: "
        + "'random' and 'degree'.",
        typeConverter=TypeConverters.toString,
    )
    srcCol: Param[str] = Param(
        Params._dummy(),
        "srcCol",
        "Name of the input column for source vertex IDs.",
        typeConverter=TypeConverters.toString,
    )
    dstCol: Param[str] = Param(
        Params._dummy(),
        "dstCol",
        "Name of the input column for destination vertex IDs.",
        typeConverter=TypeConverters.toString,
    )

    def __init__(self, *args: Any):
        super(_PowerIterationClusteringParams, self).__init__(*args)
        self._setDefault(k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst")

    @since("2.4.0")
    def getK(self) -> int:
        """
        Gets the value of :py:attr:`k` or its default value.
        """
        return self.getOrDefault(self.k)

    @since("2.4.0")
    def getInitMode(self) -> str:
        """
        Gets the value of :py:attr:`initMode` or its default value.
        """
        return self.getOrDefault(self.initMode)

    @since("2.4.0")
    def getSrcCol(self) -> str:
        """
        Gets the value of :py:attr:`srcCol` or its default value.
        """
        return self.getOrDefault(self.srcCol)

    @since("2.4.0")
    def getDstCol(self) -> str:
        """
        Gets the value of :py:attr:`dstCol` or its default value.
        """
        return self.getOrDefault(self.dstCol)


@inherit_doc
class PowerIterationClustering(
    _PowerIterationClusteringParams,
    JavaParams,
    JavaMLReadable["PowerIterationClustering"],
    JavaMLWritable,
):
    """
    Power Iteration Clustering (PIC), a scalable graph clustering algorithm developed by
    `Lin and Cohen <http://www.cs.cmu.edu/~frank/papers/icml2010-pic-final.pdf>`_. From the
    abstract: PIC finds a very low-dimensional embedding of a dataset using truncated power
    iteration on a normalized pair-wise similarity matrix of the data.

    This class is not yet an Estimator/Transformer, use :py:func:`assignClusters` method
    to run the PowerIterationClustering algorithm.

    .. versionadded:: 2.4.0

    Notes
    -----
    See `Wikipedia on Spectral clustering <http://en.wikipedia.org/wiki/Spectral_clustering>`_

    Examples
    --------
    >>> 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).toDF("src", "dst", "weight").repartition(1)
    >>> pic = PowerIterationClustering(k=2, weightCol="weight")
    >>> pic.setMaxIter(40)
    PowerIterationClustering...
    >>> assignments = pic.assignClusters(df)
    >>> assignments.sort(assignments.id).show(truncate=False)
    +---+-------+
    |id |cluster|
    +---+-------+
    |0  |0      |
    |1  |0      |
    |2  |0      |
    |3  |0      |
    |4  |0      |
    |5  |1      |
    +---+-------+
    ...
    >>> pic_path = temp_path + "/pic"
    >>> pic.save(pic_path)
    >>> pic2 = PowerIterationClustering.load(pic_path)
    >>> pic2.getK()
    2
    >>> pic2.getMaxIter()
    40
    >>> pic2.assignClusters(df).take(6) == assignments.take(6)
    True
    """

    _input_kwargs: Dict[str, Any]

    @keyword_only
    def __init__(
        self,
        *,
        k: int = 2,
        maxIter: int = 20,
        initMode: str = "random",
        srcCol: str = "src",
        dstCol: str = "dst",
        weightCol: Optional[str] = None,
    ):
        """
        __init__(self, \\*, k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst",\
                 weightCol=None)
        """
        super(PowerIterationClustering, self).__init__()
        self._java_obj = self._new_java_obj(
            "org.apache.spark.ml.clustering.PowerIterationClustering", self.uid
        )
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    @since("2.4.0")
    def setParams(
        self,
        *,
        k: int = 2,
        maxIter: int = 20,
        initMode: str = "random",
        srcCol: str = "src",
        dstCol: str = "dst",
        weightCol: Optional[str] = None,
    ) -> "PowerIterationClustering":
        """
        setParams(self, \\*, k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst",\
                  weightCol=None)
        Sets params for PowerIterationClustering.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    @since("2.4.0")
    def setK(self, value: int) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`k`.
        """
        return self._set(k=value)

    @since("2.4.0")
    def setInitMode(self, value: str) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`initMode`.
        """
        return self._set(initMode=value)

    @since("2.4.0")
    def setSrcCol(self, value: str) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`srcCol`.
        """
        return self._set(srcCol=value)

    @since("2.4.0")
    def setDstCol(self, value: str) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`dstCol`.
        """
        return self._set(dstCol=value)

    @since("2.4.0")
    def setMaxIter(self, value: int) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`maxIter`.
        """
        return self._set(maxIter=value)

    @since("2.4.0")
    def setWeightCol(self, value: str) -> "PowerIterationClustering":
        """
        Sets the value of :py:attr:`weightCol`.
        """
        return self._set(weightCol=value)

    @since("2.4.0")
    def assignClusters(self, dataset: DataFrame) -> DataFrame:
        """
        Run the PIC algorithm and returns a cluster assignment for each input vertex.

        Parameters
        ----------
        dataset : :py:class:`pyspark.sql.DataFrame`
          A dataset with columns src, dst, weight representing the affinity matrix,
          which is the matrix A in the PIC paper. Suppose the src column value is i,
          the dst column value is j, the weight column value is similarity s,,ij,,
          which must be nonnegative. This is a symmetric matrix and hence
          s,,ij,, = s,,ji,,. For any (i, j) with nonzero similarity, there should be
          either (i, j, s,,ij,,) or (j, i, s,,ji,,) in the input. Rows with i = j are
          ignored, because we assume s,,ij,, = 0.0.

        Returns
        -------
        :py:class:`pyspark.sql.DataFrame`
            A dataset that contains columns of vertex id and the corresponding cluster for
            the id. The schema of it will be:
            - id: Long
            - cluster: Int
        """
        self._transfer_params_to_java()
        assert self._java_obj is not None

        jdf = self._java_obj.assignClusters(dataset._jdf)
        return DataFrame(jdf, dataset.sparkSession)


if __name__ == "__main__":
    import doctest
    import numpy
    import pyspark.ml.clustering
    from pyspark.sql import SparkSession

    try:
        # Numpy 1.14+ changed it's string format.
        numpy.set_printoptions(legacy="1.13")
    except TypeError:
        pass
    globs = pyspark.ml.clustering.__dict__.copy()
    # The small batch size here ensures that we see multiple batches,
    # even in these small test examples:
    spark = SparkSession.builder.master("local[2]").appName("ml.clustering tests").getOrCreate()
    sc = spark.sparkContext
    globs["sc"] = sc
    globs["spark"] = spark
    import tempfile

    temp_path = tempfile.mkdtemp()
    globs["temp_path"] = temp_path
    try:
        (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS)
        spark.stop()
    finally:
        from shutil import rmtree

        try:
            rmtree(temp_path)
        except OSError:
            pass
    if failure_count:
        sys.exit(-1)