package cn.edu.whu.boot.common.utils.allocation;

import org.apache.commons.math3.ml.clustering.CentroidCluster;
import org.apache.commons.math3.ml.clustering.KMeansPlusPlusClusterer;
import org.locationtech.jts.geom.Point;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * 实现 K Means 聚类
 */
public class KMeansClustering {
    public static Map<Integer, List<Point>> clustering(List<Point> points, int k) {
        List<PointWrapper> clusterInput = points.stream()
                .map(PointWrapper::new).collect(Collectors.toList());
        KMeansPlusPlusClusterer<PointWrapper> clusterer = new KMeansPlusPlusClusterer<>(k, 10000);
        List<CentroidCluster<PointWrapper>> clusterResults = clusterer.cluster(clusterInput);
        Map<Integer, List<Point>> res = new HashMap<>();
        IntStream
                .range(0, clusterResults.size())
                .forEach(
                        i -> res.put(
                                i,
                                clusterResults.get(i).getPoints().stream().map(PointWrapper::getOriginal).collect(Collectors.toList())
                        )
                );
        return res;
    }
}

