from util.system.types cimport i64
from util.generic.string cimport TString


cdef extern from "datacloud/ml_utils/kmeans/lib/fast_kmeans.h" namespace "Datacloud::MLUtils::KMeans" nogil:
    i64 Expectation(
        const TString& ytToken, const TString& cluster, i64 nClusters,
        const TString& clusterCentersTable, const TString& objectsTable,
        const TString& outputExpectationTable) nogil

    void Minimization(
        const TString& ytToken, const TString& cluster,
        const TString& expectationTable, const TString& outputClusterCenterTable) nogil

    void ExpectationMinimization(
        const TString& ytToken, const TString& cluster, i64 nClusters,
        const TString& objectsTable, const TString& clusterCentersTable,
        const TString& outputClusterCentersTable) nogil


cdef _expectation(
    const TString& ytToken, const TString& cluster, i64 nClusters,
    const TString& clusterCentersTable, const TString& objectsTable,
    const TString& outputExpectationTable):
    cdef i64 val
    with nogil:
        val = Expectation(ytToken, cluster, nClusters, clusterCentersTable, objectsTable, outputExpectationTable)
    return val


def expectation(ytToken, cluster, nClusters, clusterCentersTable, objectsTable, outputExpectationTable):
    return _expectation(ytToken, cluster, nClusters, clusterCentersTable, objectsTable, outputExpectationTable)


cdef _minimization(
    const TString& ytToken, const TString& cluster,
    const TString& expectationTable, const TString& outputClusterCenterTable):
    with nogil:
        Minimization(ytToken, cluster, expectationTable, outputClusterCenterTable)


def minimization(ytToken, cluster, expectationTable, outputClusterCenterTable):
    _minimization(ytToken, cluster, expectationTable, outputClusterCenterTable)


cdef _expectation_minimization(
    const TString& ytToken, const TString& cluster, i64 nClusters,
    const TString& objectsTable, const TString& clusterCentersTable,
    const TString& outputClusterCentersTable):
    with nogil:
        ExpectationMinimization(ytToken, cluster, nClusters, objectsTable, clusterCentersTable, outputClusterCentersTable)


def expectation_minimization(ytToken, cluster, nClusters, objectsTable, expectationTable, outputClusterCentersTable):
    _expectation_minimization(ytToken, cluster, nClusters, objectsTable, expectationTable, outputClusterCentersTable)
