1 files changed, 11 insertions, 4 deletions
diff --git a/ml/kmeans/KMeans.h b/ml/KMeans.h
index 4ea3b6a89..0398eeb86 100644
--- a/ml/kmeans/KMeans.h
+++ b/ml/KMeans.h
@@ -9,6 +9,7 @@
 #include <mutex>
 
 #include "SamplesBuffer.h"
+#include "json/single_include/nlohmann/json.hpp"
 
 class KMeans {
 public:
@@ -17,8 +18,16 @@ public:
         MaxDist = std::numeric_limits<CalculatedNumber>::min();
     };
 
-    void train(SamplesBuffer &SB, size_t MaxIterations);
-    CalculatedNumber anomalyScore(SamplesBuffer &SB);
+    void train(const std::vector<DSample> &Samples, size_t MaxIterations);
+    CalculatedNumber anomalyScore(const DSample &Sample) const;
+
+    void toJson(nlohmann::json &J) const {
+        J = nlohmann::json{
+            {"CCs", ClusterCenters},
+            {"MinDist", MinDist},
+            {"MaxDist", MaxDist}
+        };
+    }
 
 private:
     size_t NumClusters;
@@ -27,8 +36,6 @@ private:
 
     CalculatedNumber MinDist;
     CalculatedNumber MaxDist;
-
-    std::mutex Mutex;
 };
 
 #endif /* KMEANS_H */