from util.generic.string cimport TString
from util.generic.vector cimport TVector


cdef extern from "datacloud/features/dssm/fast_dssm45.h" namespace "Datacloud::Features::DSSM" nogil:
    void FastDSSMStep5(
        const TString&, const TString&, const TString&, const TString&) nogil
    void FastUniqueHashReduce(
        const TString&, const TString&, const TString&, const TString&) nogil
    void FastJoinCidsReduce(
        const TString&, const TString&, const TString&, const TString&,
        const TString&) nogil
    void FastDotProduct(
        const TString& ytToken, const TString& cluster,
        const TVector[TVector[float]]& clusterCenters,
        const TString& inputTable, const TString& outputTable) nogil


cdef _fast_dssm_step5(
    const TString& ytToken, const TString& cluster,
    const TString& inputTable, const TString& outputTable):
    with nogil:
        FastDSSMStep5(ytToken, cluster, inputTable, outputTable)


def fast_dssm_step5(ytToken, cluster, inputTable, outputTable):
    _fast_dssm_step5(ytToken, cluster, inputTable, outputTable)


cdef _fast_unique_hash_reduce(
    const TString& ytToken, const TString& cluster,
    const TString& inputTable, const TString& outputTable):
    with nogil:
        FastUniqueHashReduce(ytToken, cluster, inputTable, outputTable)


def fast_unique_hash_reduce(ytToken, cluster, inputTable, outputTable):
    _fast_unique_hash_reduce(ytToken, cluster, inputTable, outputTable)


cdef _fast_join_cids_reduce(
    const TString& ytToken, const TString& cluster, const TString& inputDotTable,
    const TString& inputHash2KeyTable, const TString& outputTable):
    with nogil:
        FastJoinCidsReduce(ytToken, cluster, inputDotTable, inputHash2KeyTable, outputTable)


def fast_join_cids_reduce(ytToken, cluster, inputDotTable, inputHash2KeyTable,
                          outputTable):
    _fast_join_cids_reduce(ytToken, cluster, inputDotTable, inputHash2KeyTable,
                           outputTable)


cdef _fast_dot_product(
    const TString& ytToken, const TString& cluster,
    const TVector[TVector[float]]& clusterCenters,
    const TString& inputTable, const TString& outputTable):
    with nogil:
        FastDotProduct(ytToken, cluster, clusterCenters,
                       inputTable, outputTable)


def fast_dot_product(ytToken, cluster, clusterCenters, inputTable, outputTable):
    _fast_dot_product(ytToken, cluster, clusterCenters, inputTable, outputTable)
