package cn.edu.bjut.textmining.chapter5;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

import cn.edu.bjut.textmining.util.TextVectorizer;

/** 球面K-Means：单位向量按余弦分配，簇向量之和重新单位化。 */
public final class KMeansEngine {
    private KMeansEngine() {}

    public static final class Result {
        private final int[] labels;
        private final int iterations;
        private final double[][] centers;
        private final boolean converged;

        public Result(int[] labels, int iterations, double[][] centers, boolean converged) {
            this.labels = labels;
            this.iterations = iterations;
            this.centers = centers;
            this.converged = converged;
        }

        public int[] labels() {
            return labels.clone();
        }

        public int iterations() {
            return iterations;
        }

        public double[][] centers() {
            return centers;
        }

        public boolean converged() {
            return converged;
        }
    }

    public static Result fit(double[][] data, int k, int maxIterations, int... initial) {
        validate(data);
        if (k < 1 || k > data.length || initial.length != k || maxIterations < 1) {
            throw new IllegalArgumentException("请检查簇数、初始中心数和迭代上限");
        }
        int[] labels = new int[data.length];
        Arrays.fill(labels, -1);
        double[][] points = new double[data.length][];
        for (int i = 0; i < data.length; i++) {
            points[i] = data[i].clone();
            TextVectorizer.normalize(points[i]);
        }
        double[][] centers = new double[k][];
        Set<Integer> seen = new HashSet<>();
        for (int c = 0; c < k; c++) {
            if (initial[c] < 0 || initial[c] >= data.length || !seen.add(initial[c])) {
                throw new IllegalArgumentException("初始中心索引必须有效且互不相同");
            }
            centers[c] = points[initial[c]].clone();
        }
        int completed = 0;
        boolean converged = false;
        for (int iteration = 0; iteration < maxIterations; iteration++) {
            completed++;
            boolean changed = false;
            for (int i = 0; i < points.length; i++) {
                int best = 0;
                double score = -Double.MAX_VALUE;
                for (int c = 0; c < k; c++) {
                    double candidate = TextVectorizer.cosine(points[i], centers[c]);
                    if (candidate > score) {
                        score = candidate;
                        best = c;
                    }
                }
                if (labels[i] != best) {
                    labels[i] = best;
                    changed = true;
                }
            }
            // 空簇保留上一轮中心；不把零向量当作有效中心。
            for (int c = 0; c < k; c++) {
                double[] sum = TextVectorizer.centroid(points, labels, c);
                if (squaredNorm(sum) > 0.0) {
                    centers[c] = sum;
                }
            }
            if (!changed) {
                converged = true;
                break;
            }
        }
        return new Result(labels, completed, centers, converged);
    }

    public static void validate(double[][] points) {
        if (points == null || points.length == 0 || points[0].length == 0) {
            throw new IllegalArgumentException("输入矩阵不能为空");
        }
        int d = points[0].length;
        for (double[] row : points) {
            if (row.length != d) {
                throw new IllegalArgumentException("向量维度必须相同");
            }
            for (double value : row) {
                if (!Double.isFinite(value)) {
                    throw new IllegalArgumentException("向量必须是有限数值");
                }
            }
            if (squaredNorm(row) == 0.0) {
                throw new IllegalArgumentException("零向量没有余弦方向，请先移除");
            }
        }
    }

    public static double squaredNorm(double[] row) {
        double value = 0;
        for (double x : row) {
            value += x * x;
        }
        return value;
    }

    public static int[] farthestFirst(double[][] points, int k) {
        validate(points);
        if (k < 1 || k > points.length) {
            throw new IllegalArgumentException("簇数无效");
        }
        int[] selected = new int[k];
        boolean[] used = new boolean[points.length];
        used[0] = true;
        for (int c = 1; c < k; c++) {
            double best = -1;
            int index = -1;
            for (int i = 0; i < points.length; i++) {
                if (used[i]) {
                    continue;
                }
                double nearest = Double.POSITIVE_INFINITY;
                for (int j = 0; j < c; j++) {
                    nearest =
                            Math.min(
                                    nearest,
                                    TextVectorizer.cosineDistance(points[i], points[selected[j]]));
                }
                if (nearest > best) {
                    best = nearest;
                    index = i;
                }
            }
            selected[c] = index;
            used[index] = true;
        }
        return selected;
    }

    public static double objective(double[][] points, Result result) {
        double sum = 0;
        for (int i = 0; i < points.length; i++) {
            sum += TextVectorizer.cosineDistance(points[i], result.centers[result.labels[i]]);
        }
        return sum;
    }
}