/*
 * Copyright 2021 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * ---------------------------------------------------------------------
 *
 * Portions of this file were modified and ported to Vala by Victor Vintorez in 2026.
 * These modifications are licensed under the European Union Public License (EUPL),
 * Version 1.2 or later. You may obtain a copy of the License at:
 * https://joinup.ec.europa.eu/collection/eupl/eupl-text-eupl-12
 */

using GLib;

namespace MaterialColorUtilities {
	/**
	 * An image quantizer that improves on the speed of a standard K-Means algorithm by implementing
	 * several optimizations, including deduping identical pixels and a triangle inequality rule that
	 * reduces the number of comparisons needed to identify which cluster a point should be moved to.
	 *
	 * <p>Wsmeans stands for Weighted Square Means.
	 *
	 * <p>This algorithm was designed by M. Emre Celebi, and was found in their 2011 paper, Improving
	 * the Performance of K-Means for Color Quantization. https://arxiv.org/abs/1101.0395
	 */
	public sealed class QuantizerWsmeans : Object {
		private QuantizerWsmeans() {}

		private struct Distance {
			int index;
			double distance;

			public static int compare_to(ref Distance a, ref Distance b) {
				if (a.distance < b.distance)return -1;
				if (a.distance > b.distance)return 1;
				return 0;
			}
		}

		private const int MAX_ITERATIONS = 10;
		private const double MIN_MOVEMENT_DISTANCE = 3.0;

		/**
		 * Reduce the number of colors needed to represented the input, minimizing the difference between
		 * the original image and the recolored image.
		 *
		 * @param inputPixels Colors in ARGB format.
		 * @param startingClusters Defines the initial state of the quantizer. Passing an empty array is
		 *     fine, the implementation will create its own initial state that leads to reproducible
		 *     results for the same inputs. Passing an array that is the result of Wu quantization leads
		 *     to higher quality results.
		 * @param maxColors The number of colors to divide the image into. A lower number of colors may be
		 *     returned.
		 * @return Map with keys of colors in ARGB format, values of how many of the input pixels belong
		 *     to the color.
		 */
		public static HashTable<int, int> quantize(int[] input_pixels, int[] starting_clusters, int max_colors) {
			// Uses a seeded random number generator to ensure consistent results.
			var random = new Rand.with_seed(0x42688);

			HashTable<int, int> pixel_to_count = new HashTable<int, int> (direct_hash, direct_equal);
			double[,] points = new double[input_pixels.length, 3];
			int[] pixels = new int[input_pixels.length];
			PointProvider point_provider = new PointProviderLab();

			int point_count = 0;
			for (int i = 0; i < input_pixels.length; i++) {
				int input_pixel = input_pixels[i];
				if (!pixel_to_count.contains(input_pixel)) {
					double[] lab = point_provider.from_int(input_pixel);
					points[point_count, 0] = lab[0];
					points[point_count, 1] = lab[1];
					points[point_count, 2] = lab[2];
					pixels[point_count] = input_pixel;
					point_count++;

					pixel_to_count.insert(input_pixel, 1);
				} else {
					pixel_to_count.insert(input_pixel, pixel_to_count.get(input_pixel) + 1);
				}
			}

			int[] counts = new int[point_count];
			for (int i = 0; i < point_count; i++) {
				int pixel = pixels[i];
				int count = pixel_to_count.get(pixel);
				counts[i] = count;
			}

			int cluster_count = (int) Math.fmin(max_colors, point_count);
			if (starting_clusters.length != 0) {
				cluster_count = (int) Math.fmin(cluster_count, starting_clusters.length);
			}

			double[,] clusters = new double[cluster_count, 3];
			int clusters_created = 0;
			for (int i = 0; i < starting_clusters.length; i++) {
				var pt = point_provider.from_int(starting_clusters[i]);
				clusters[i, 0] = pt[0];
				clusters[i, 1] = pt[1];
				clusters[i, 2] = pt[2];
				clusters_created++;
			}

			int additional_clusters_needed = cluster_count - clusters_created;
			if (additional_clusters_needed > 0) {
				for (int i = 0; i < additional_clusters_needed; i++) {}
			}

			int[] cluster_indices = new int[point_count];
			for (int i = 0; i < point_count; i++) {
				cluster_indices[i] = random.int_range(0, cluster_count);
			}

			int[] index_matrix = new int[cluster_count * cluster_count];

			Distance[] distance_to_index_matrix = new Distance[cluster_count * cluster_count];
			for (int i = 0; i < cluster_count * cluster_count; i++) {
				distance_to_index_matrix[i] = Distance();
			}

			int[] pixel_count_sums = new int[cluster_count];
			for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
				for (int i = 0; i < cluster_count; i++) {
					double[] cluster_i = { clusters[i, 0], clusters[i, 1], clusters[i, 2] };
					for (int j = i + 1; j < cluster_count; j++) {
						double[] cluster_j = { clusters[j, 0], clusters[j, 1], clusters[j, 2] };
						double distance = point_provider.distance(cluster_i, cluster_j);
						distance_to_index_matrix[j * cluster_count + i].distance = distance;
						distance_to_index_matrix[j * cluster_count + i].index = i;
						distance_to_index_matrix[i * cluster_count + j].distance = distance;
						distance_to_index_matrix[i * cluster_count + j].index = j;
					}
					Distance[] row = distance_to_index_matrix[i * cluster_count : (i + 1) * cluster_count];
					for (int s = 1; s < row.length; s++) {
						Distance key = row[s];
						int t = s - 1;
						while (t >= 0 && Distance.compare_to(ref row[t], ref key) > 0) {
							row[t + 1] = row[t];
							t--;
						}
						row[t + 1] = key;
					}
					for (int j = 0; j < cluster_count; j++) {
						distance_to_index_matrix[i * cluster_count + j] = row[j];
						index_matrix[i * cluster_count + j] = row[j].index;
					}
				}

				int points_moved = 0;
				for (int i = 0; i < point_count; i++) {
					double[] point = { points[i, 0], points[i, 1], points[i, 2] };
					int previous_cluster_index = cluster_indices[i];
					double[] previous_cluster = { clusters[previous_cluster_index, 0], clusters[previous_cluster_index, 1], clusters[previous_cluster_index, 2] };
					double previous_distance = point_provider.distance(point, previous_cluster);

					double minimum_distance = previous_distance;
					int new_cluster_index = -1;
					for (int j = 0; j < cluster_count; j++) {
						if (distance_to_index_matrix[previous_cluster_index * cluster_count + j].distance >= 4 * previous_distance) {
							continue;
						}
						double[] cluster_j = { clusters[j, 0], clusters[j, 1], clusters[j, 2] };
						double distance = point_provider.distance(point, cluster_j);
						if (distance < minimum_distance) {
							minimum_distance = distance;
							new_cluster_index = j;
						}
					}
					if (new_cluster_index != -1) {
						double distance_change = Math.fabs(Math.sqrt(minimum_distance) - Math.sqrt(previous_distance));
						if (distance_change > MIN_MOVEMENT_DISTANCE) {
							points_moved++;
							cluster_indices[i] = new_cluster_index;
						}
					}
				}

				if (points_moved == 0 && iteration != 0) {
					break;
				}

				double[] component_a_sums = new double[cluster_count];
				double[] component_b_sums = new double[cluster_count];
				double[] component_c_sums = new double[cluster_count];
				for (int j = 0; j < pixel_count_sums.length; j++) {
					pixel_count_sums[j] = 0;
				}
				for (int i = 0; i < point_count; i++) {
					int cluster_index = cluster_indices[i];
					double[] point = { points[i, 0], points[i, 1], points[i, 2] };
					int count = counts[i];
					pixel_count_sums[cluster_index] += count;
					component_a_sums[cluster_index] += (point[0] * count);
					component_b_sums[cluster_index] += (point[1] * count);
					component_c_sums[cluster_index] += (point[2] * count);
				}

				for (int i = 0; i < cluster_count; i++) {
					int count = pixel_count_sums[i];
					if (count == 0) {
						clusters[i, 0] = 0.0;
						clusters[i, 1] = 0.0;
						clusters[i, 2] = 0.0;
						continue;
					}
					double a = component_a_sums[i] / count;
					double b = component_b_sums[i] / count;
					double c = component_c_sums[i] / count;
					clusters[i, 0] = a;
					clusters[i, 1] = b;
					clusters[i, 2] = c;
				}
			}

			HashTable<int, int> argb_to_population = new HashTable<int, int> (direct_hash, direct_equal);
			for (int i = 0; i < cluster_count; i++) {
				int count = pixel_count_sums[i];
				if (count == 0) {
					continue;
				}

				double[] cluster_i = { clusters[i, 0], clusters[i, 1], clusters[i, 2] };
				int possible_new_cluster = point_provider.to_int(cluster_i);
				if (argb_to_population.contains(possible_new_cluster)) {
					continue;
				}

				argb_to_population.insert(possible_new_cluster, count);
			}

			return argb_to_population;
		}
	}
}
