package at.procon.dip.clustering.spi;

import at.procon.dip.clustering.ClusteringAlgorithm;
import at.procon.dip.clustering.dto.ClusteringEngineAssignment;
import at.procon.dip.clustering.dto.ClusteringEngineCluster;
import at.procon.dip.clustering.dto.ClusteringEngineRequest;
import at.procon.dip.clustering.dto.ClusteringEngineResult;
import at.procon.dip.clustering.dto.SelectedEmbeddingRow;
import java.util.ArrayList;
import java.util.List;
import org.apache.commons.math3.ml.clustering.CentroidCluster;
import org.apache.commons.math3.ml.clustering.DoublePoint;
import org.apache.commons.math3.ml.clustering.KMeansPlusPlusClusterer;
import org.springframework.stereotype.Component;

@Component
public class KMeansClusteringEngine implements EmbeddingClusteringEngine {

    @Override
    public ClusteringAlgorithm algorithm() {
        return ClusteringAlgorithm.KMEANS;
    }

    @Override
    public ClusteringEngineResult cluster(List<SelectedEmbeddingRow> items, ClusteringEngineRequest request) {
        if (items == null || items.isEmpty()) {
            throw new IllegalArgumentException("Selection contains no embeddings to cluster");
        }
        if (request == null) {
            throw new IllegalArgumentException("Missing clustering request");
        }

        int k = request.requiredInt("k");
        if (k <= 0) {
            throw new IllegalArgumentException("KMeans requires k > 0");
        }
        if (k > items.size()) {
            throw new IllegalArgumentException("KMeans k must be <= selected item count");
        }

        int maxIterations = request.intValue("maxIterations", 100);

        List<IndexedPoint> points = new ArrayList<>(items.size());
        for (int i = 0; i < items.size(); i++) {
            points.add(new IndexedPoint(i, toDouble(items.get(i).embeddingVector())));
        }

        KMeansPlusPlusClusterer<IndexedPoint> clusterer = new KMeansPlusPlusClusterer<>(k, maxIterations);
        List<CentroidCluster<IndexedPoint>> clusters = clusterer.cluster(points);

        List<ClusteringEngineCluster> resultClusters = new ArrayList<>();
        List<ClusteringEngineAssignment> assignments = new ArrayList<>();

        for (int label = 0; label < clusters.size(); label++) {
            CentroidCluster<IndexedPoint> cluster = clusters.get(label);
            resultClusters.add(new ClusteringEngineCluster(label, cluster.getPoints().size(), false));
            double[] centroid = cluster.getCenter().getPoint();
            for (IndexedPoint point : cluster.getPoints()) {
                SelectedEmbeddingRow item = items.get(point.index());
                assignments.add(new ClusteringEngineAssignment(
                        item.embeddingId(),
                        item.documentId(),
                        item.representationId(),
                        label,
                        euclidean(point.getPoint(), centroid),
                        null,
                        false
                ));
            }
        }

        return new ClusteringEngineResult(resultClusters, assignments, 0L);
    }

    private double[] toDouble(float[] values) {
        double[] result = new double[values.length];
        for (int i = 0; i < values.length; i++) {
            result[i] = values[i];
        }
        return result;
    }

    private double euclidean(double[] a, double[] b) {
        double sum = 0.0d;
        for (int i = 0; i < a.length; i++) {
            double d = a[i] - b[i];
            sum += d * d;
        }
        return Math.sqrt(sum);
    }

    private static class IndexedPoint extends DoublePoint {
        private final int index;

        IndexedPoint(int index, double[] point) {
            super(point);
            this.index = index;
        }

        int index() {
            return index;
        }
    }
}
