package cn.edu.bjut.textmining.chapter5;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/** s(i,k)=-||x_i-x_k||²；责任度和可用度采用同步更新与阻尼。 */
public final class AffinityPropagation {
    private AffinityPropagation() {}

    public static final class Result {
        public final int[] labels;
        public final int[] exemplars;
        public final int iterations;
        public final boolean converged;
        public final double preference;

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

    public static double[][] similarities(double[][] vectors) {
        KMeansEngine.validate(vectors);
        double[][] s = new double[vectors.length][vectors.length];
        for (int i = 0; i < s.length; i++) {
            for (int k = 0; k < s.length; k++) {
                for (int j = 0; j < vectors[i].length; j++) {
                    double delta = vectors[i][j] - vectors[k][j];
                    s[i][k] -= delta * delta;
                }
            }
        }
        return s;
    }

    public static double medianPreference(double[][] s) {
        if (s.length < 2) {
            throw new IllegalArgumentException("AP至少需要两个样本");
        }
        double[] values = new double[s.length * (s.length - 1)];
        int p = 0;
        for (int i = 0; i < s.length; i++) {
            for (int k = 0; k < s.length; k++) {
                if (i != k) {
                    values[p++] = s[i][k];
                }
            }
        }
        Arrays.sort(values);
        return (values[(p - 1) / 2] + values[p / 2]) / 2;
    }

    public static Result fit(
            double[][] vectors, double damping, int maxIterations, int stableRounds) {
        double[][] s = similarities(vectors);
        int n = s.length;
        double preference = medianPreference(s);
        if (damping < 0.5 || damping >= 1 || maxIterations < 1 || stableRounds < 1) {
            throw new IllegalArgumentException("请检查阻尼、迭代上限和稳定轮数");
        }
        for (int i = 0; i < n; i++) {
            s[i][i] = preference;
        }
        double[][] r = new double[n][n], a = new double[n][n];
        boolean[] previous = new boolean[n];
        int stable = 0, completed = 0;
        boolean converged = false;
        for (int iteration = 0; iteration < maxIterations; iteration++) {
            completed++;
            updateResponsibilities(s, r, a, damping);
            updateAvailabilities(r, a, damping);
            boolean[] current = new boolean[n];
            int count = 0;
            for (int k = 0; k < n; k++) {
                current[k] = r[k][k] + a[k][k] > 0;
                if (current[k]) {
                    count++;
                }
            }
            stable = Arrays.equals(previous, current) ? stable + 1 : 0;
            previous = current;
            if (count > 0 && stable >= stableRounds) {
                converged = true;
                break;
            }
        }
        List<Integer> chosen = new ArrayList<>();
        for (int k = 0; k < n; k++) {
            if (r[k][k] + a[k][k] > 0) {
                chosen.add(k);
            }
        }
        if (chosen.isEmpty()) {
            throw new IllegalStateException("AP尚未产生代表点，请调整偏好、阻尼或迭代上限");
        }
        int[] labels = new int[n], exemplars = new int[chosen.size()];
        for (int c = 0; c < exemplars.length; c++) {
            exemplars[c] = chosen.get(c);
        }
        // 以原始相似度分配给代表点；代表点归入自身簇。
        for (int i = 0; i < n; i++) {
            double best = -Double.MAX_VALUE;
            for (int c = 0; c < exemplars.length; c++) {
                if (s[i][exemplars[c]] > best) {
                    best = s[i][exemplars[c]];
                    labels[i] = c;
                }
            }
        }
        for (int c = 0; c < exemplars.length; c++) {
            labels[exemplars[c]] = c;
        }
        return new Result(labels, exemplars, completed, converged, preference);
    }

    public static void updateResponsibilities(
            double[][] s, double[][] r, double[][] a, double damping) {
        for (int i = 0; i < s.length; i++) {
            for (int k = 0; k < s.length; k++) {
                double maxOther = -Double.MAX_VALUE;
                for (int j = 0; j < s.length; j++) {
                    if (j != k) {
                        maxOther = Math.max(maxOther, a[i][j] + s[i][j]);
                    }
                }
                double next = s[i][k] - maxOther;
                r[i][k] = damping * r[i][k] + (1 - damping) * next;
            }
        }
    }

    public static void updateAvailabilities(double[][] r, double[][] a, double damping) {
        for (int k = 0; k < r.length; k++) {
            double positiveSum = 0;
            for (int i = 0; i < r.length; i++) {
                if (i != k) {
                    positiveSum += Math.max(0, r[i][k]);
                }
            }
            for (int i = 0; i < r.length; i++) {
                double next =
                        i == k
                                ? positiveSum
                                : Math.min(0, r[k][k] + positiveSum - Math.max(0, r[i][k]));
                a[i][k] = damping * a[i][k] + (1 - damping) * next;
            }
        }
    }
}