package com.one;

import com.vmx.ProgressCallback;
import javax.microedition.lcdui.Image;

/**
 * An efficient color quantization algorithm, adapted from the C++
 * implementation quantize.c in <a href="http://www.imagemagick.org/">ImageMagick</a>. The pixels for
 * an image are placed into an oct tree. The oct tree is reduced in
 * size, and the pixels from the original image are reassigned to the
 * nodes in the reduced tree.<p>
 *
 * @version 0.90 19 Sep 2000
 * @author <a href="http://www.gurge.com/amd/">Adam Doppelt</a>
 */
public class ColorQuantizer
{
	public static final int MAX_NODES = 266817;
	public static final int MAX_NODES_ALPHA = MAX_NODES + MAX_NODES / 3;

	static final int MAX_RGB = 255;
	static final int MAX_TREE_DEPTH = 8;

	// these are precomputed in advance
	static int SQUARES[];
	static int SHIFT[];

	static
	{
		SQUARES = new int[MAX_RGB + MAX_RGB + 1];
		
		for(int i = -MAX_RGB; i <= MAX_RGB; i++)
		{
			SQUARES[i + MAX_RGB] = i * i;
		}

		SHIFT = new int[MAX_TREE_DEPTH + 1];
		
		for(int i = 0; i < MAX_TREE_DEPTH + 1; ++i)
		{
			SHIFT[i] = 1 << (15 - i);
		}
	}
	
	Cube cube;

	/**
	 * Reduce the image to the given number of colors.
	 *
	 * @param image image to be reduced
	 * @param max_colors maximum number of colors the image may consist of
	 * @param max_nodes maximum number of unique colors simultaneously existing in an octree
	 * @param accurate actually serach for the best plette entry for the given color
	 * @param usealpha keep alpha values of the image in account (use hextree instead of octree)
	 * @param callback for the progress tracking
	 */
	public ColorQuantizer(Image image, int max_colors, int max_nodes, boolean accurate, boolean usealpha, ProgressCallback callback)
	{
		cube = new Cube(image, max_colors, max_nodes, accurate, usealpha, callback);
		
		cube.classification();
		cube.reduction();
		cube.assignment();
	}

	/**
	 * Create ColorQuantizer from existing palette.
	 * Most useful for inacurate but fast color indexing.
	 *
	 * @param palette the palete to create the quantizer from
	 * @param accurate actually serach for the best plette entry for the given color
	 * @param usealpha keep alpha values of the image in account (use hextree instead of octree)
	 * @param callback for the progress tracking
	 */
	public ColorQuantizer(int[] palette, boolean accurate, boolean usealpha, ProgressCallback callback)
	{
		this(Image.createRGBImage(palette, palette.length, 1, true), palette.length, palette.length * 2, accurate, usealpha, callback);
	}
	
	public int[] getPalette()
	{
		return cube.colormap;
	}
	
	public int getColorIndex(int color)
	{
		return cube.colorIndex(color);
	}

	static class Cube
	{
		Image image;
		int max_colors;
		int max_nodes;
		int colormap[];

		boolean accurate;
		boolean usealpha;

		ProgressCallback callback;

		Node root;
		int depth;

		// counter for the number of colors in the cube. this gets
		// recalculated often.
		int colors;

		// counter for the number of nodes in the tree
		int nodes;

		Cube(Image image, int max_colors, int max_nodes, boolean accurate, boolean usealpha, ProgressCallback callback)
		{
			this.image = image;
			this.max_colors = max_colors;
			this.max_nodes = max_nodes;

			this.accurate = accurate;
			this.usealpha = usealpha;
			
			this.callback = callback;

			int i = max_colors;
			
			// tree_depth = log max_colors
			//                 4
			
			for(depth = 1; i != 0; depth++)
			{
				i /= 4;
			}
			
			if(depth > 1)
			{
				--depth;
			}
			
			if(depth > MAX_TREE_DEPTH)
			{
				depth = MAX_TREE_DEPTH;
			}
			else if(depth < 2)
			{
				depth = 2;
			}

			root = new Node(this);
		}

		/*
         * Procedure Classification begins by initializing a color
         * description tree of sufficient depth to represent each
         * possible input color in a leaf. However, it is impractical
         * to generate a fully-formed color description tree in the
         * classification phase for realistic values of cmax. If
         * colors components in the input image are quantized to k-bit
         * precision, so that cmax= 2k-1, the tree would need k levels
         * below the root node to allow representing each possible
         * input color in a leaf. This becomes prohibitive because the
         * tree's total number of nodes is 1 + sum(i=1,k,8k).
         *
         * A complete tree would require 19,173,961 nodes for k = 8,
		 * cmax = 255. Therefore, to avoid building a fully populated
		 * tree, QUANTIZE: (1) Initializes data structures for nodes
		 * only as they are needed; (2) Chooses a maximum depth for
		 * the tree as a function of the desired number of colors in
		 * the output image (currently log2(colormap size)).
		 *
		 * For each pixel in the input image, classification scans
		 * downward from the root of the color description tree. At
		 * each level of the tree it identifies the single node which
		 * represents a cube in RGB space containing It updates the
		 * following data for each such node:
		 *
		 *   number_pixels : Number of pixels whose color is contained
		 *   in the RGB cube which this node represents;
		 *
		 *   unique : Number of pixels whose color is not represented
		 *   in a node at lower depth in the tree; initially, n2 = 0
		 *   for all nodes except leaves of the tree.
		 *
		 *   total_red/green/blue : Sums of the red, green, and blue
		 *   component values for all pixels not classified at a lower
		 *   depth. The combination of these sums and n2 will
		 *   ultimately characterize the mean color of a set of pixels
		 *   represented by this node.
		 */
		void classification()
		{
			int width = image.getWidth();
			int height = image.getHeight();
			
			int[] pixel = new int[width];

			if(callback != null)
			{
				callback.setProgress(0);
				callback.setMax(height);
			}
	
			for(int y = 0; y < height; y++)
			{
				image.getRGB(pixel, 0, width, 0, y, width, 1);

				for(int x = 0; x < width; x++)
				{
					int red = (pixel[x] >> 16) & 0xFF;
					int green = (pixel[x] >> 8) & 0xFF;
					int blue = (pixel[x] >> 0) & 0xFF;
					int alpha = (pixel[x] >> 24) & 0xFF;

					// a hard limit on the number of nodes in the tree
					if(nodes > max_nodes)
					{
						root.pruneLevel();
						--depth;
					}

					// walk the tree to depth, increasing the
					// number_pixels count for each node
					Node node = root;

					if(usealpha)
					{
						for(int level = 1; level <= depth; ++level)
						{
							int id = (((red > node.mid_red ? 1 : 0) << 0) |
									  ((green > node.mid_green ? 1 : 0) << 1) |
									  ((blue > node.mid_blue ? 1 : 0) << 2) |
									  ((alpha > node.mid_alpha ? 1 : 0) << 3));

							if(node.child[id] == null)
							{
								new Node(node, id, level);
							}

							node = node.child[id];
							node.number_pixels += SHIFT[level];
						}
					}
					else
					{
						for(int level = 1; level <= depth; ++level)
						{
							int id = (((red > node.mid_red ? 1 : 0) << 0) |
									  ((green > node.mid_green ? 1 : 0) << 1) |
									  ((blue > node.mid_blue ? 1 : 0) << 2));

							if(node.child[id] == null)
							{
								new Node(node, id, level);
							}

							node = node.child[id];
							node.number_pixels += SHIFT[level];
						}
					}

					++node.unique;

					node.total_red += red;
					node.total_green += green;
					node.total_blue += blue;
					node.total_alpha += alpha;
				}

				if(callback != null)
				{
					callback.progress(1);
				}
			}
		}

		/*
         * reduction repeatedly prunes the tree until the number of
         * nodes with unique > 0 is less than or equal to the maximum
         * number of colors allowed in the output image.
         *
         * When a node to be pruned has offspring, the pruning
         * procedure invokes itself recursively in order to prune the
         * tree from the leaves upward.  The statistics of the node
         * being pruned are always added to the corresponding data in
         * that node's parent.  This retains the pruned node's color
         * characteristics for later averaging.
         */
		void reduction()
		{
			long threshold = 1;
			
			while(colors > max_colors)
			{
				colors = 0;
				threshold = root.reduce(threshold, Integer.MAX_VALUE);
			}
		}

		/**
		 * The result of a closest color search.
		 */
		static class Search
		{
			int distance;
			int color_number;
		}

		/*
		 * Procedure assignment generates the output image from the
         * pruned tree. The output image consists of two parts: (1) A
         * color map, which is an array of color descriptions (RGB
         * triples) for each color present in the output image; (2) A
         * pixel array, which represents each pixel as an index into
         * the color map array.
         *
         * First, the assignment phase makes one pass over the pruned
         * color description tree to establish the image's color map.
         * For each node with n2 > 0, it divides Sr, Sg, and Sb by n2.
         * This produces the mean color of all pixels that classify no
         * lower than this node. Each of these colors becomes an entry
		 * in the color map.
         */
		void assignment()
		{
			colormap = new int[colors];

			colors = 0;
			root.colormap();
		}

		/**
         * Reclassify pixel in the pruned tree to identify
         * the deepest node containing the pixel's color.
         * @return Index in color map for given pixel.
         */
		int colorIndex(int pixel)
		{
			// convert to indexed color
			int red = (pixel >> 16) & 0xFF;
			int green = (pixel >> 8) & 0xFF;
			int blue = (pixel >> 0) & 0xFF;
			int alpha = (pixel >> 24) & 0xFF;

			// walk the tree to find the cube containing that color
			Node node = root;

			if(usealpha)
			{
				for(; ; )
				{
					int id = (((red > node.mid_red ? 1 : 0) << 0) |
							  ((green > node.mid_green ? 1 : 0) << 1) |
							  ((blue > node.mid_blue ? 1 : 0) << 2) |
							  ((alpha > node.mid_alpha ? 1 : 0) << 3));

					if(node.child[id] == null)
					{
						break;
					}

					node = node.child[id];
				}
			}
			else
			{
				for(; ; )
				{
					int id = (((red > node.mid_red ? 1 : 0) << 0) |
							  ((green > node.mid_green ? 1 : 0) << 1) |
							  ((blue > node.mid_blue ? 1 : 0) << 2));

					if(node.child[id] == null)
					{
						break;
					}

					node = node.child[id];
				}
			}

			if(accurate)
			{
				// Find the closest color.
				Search search = new Search();
				search.distance = Integer.MAX_VALUE;
				node.parent.closestColor(red, green, blue, alpha, search);
				return search.color_number;
				
			}
			else
			{
				// if QUICK is set, just use that
				// node. Strictly speaking, this isn't
				// necessarily best match.
				return node.color_number;
			}
		}

		/**
         * A single Node in the tree.
         */
		static class Node
		{
			Cube cube;

			// parent node
			Node parent;

			// child nodes
			Node child[];
			int nchild;

			// our index within our parent
			int id;
			
			// our level within the tree
			int level;
			
			// our color midpoint
			int mid_red;
			int mid_green;
			int mid_blue;
			int mid_alpha;

			// the pixel count for this node and all children
			long number_pixels;

			// the pixel count for this node
			int unique;
			
			// the sum of all pixels contained in this node
			int total_red;
			int total_green;
			int total_blue;
			int total_alpha;

			// used to build the colormap
			int color_number;

			Node(Cube cube)
			{
				this.cube = cube;
				this.parent = this;
				this.child = new Node[cube.usealpha ? 16 : 8];
				this.id = 0;
				this.level = 0;

				this.number_pixels = Integer.MAX_VALUE;

				this.mid_red = (MAX_RGB + 1) >> 1;
				this.mid_green = (MAX_RGB + 1) >> 1;
				this.mid_blue = (MAX_RGB + 1) >> 1;
				this.mid_alpha = (MAX_RGB + 1) >> 1;
			}

			Node(Node parent, int id, int level)
			{
				this.cube = parent.cube;
				this.parent = parent;
				this.child = new Node[cube.usealpha ? 16 : 8];
				this.id = id;
				this.level = level;

				// add to the cube
				++cube.nodes;
				
				if(level == cube.depth)
				{
					++cube.colors;
				}

				// add to the parent
				++parent.nchild;
				parent.child[id] = this;

				// figure out our midpoint
				int bi = (1 << (MAX_TREE_DEPTH - level)) >> 1;
				mid_red = parent.mid_red + ((id & 1) > 0 ? bi : -bi);
				mid_green = parent.mid_green + ((id & 2) > 0 ? bi : -bi);
				mid_blue = parent.mid_blue + ((id & 4) > 0 ? bi : -bi);
				mid_alpha = parent.mid_alpha + ((id & 8) > 0 ? bi : -bi);
			}

			/**
             * Remove this child node, and make sure our parent
             * absorbs our pixel statistics.
             */
			void pruneChild()
			{
				--parent.nchild;
				parent.unique += unique;
				parent.total_red += total_red;
				parent.total_green += total_green;
				parent.total_blue += total_blue;
				parent.total_alpha += total_alpha;
				parent.child[id] = null;
				--cube.nodes;
				cube = null;
				parent = null;
			}

			/**
             * Prune the lowest layer of the tree.
             */
			void pruneLevel()
			{
				if(nchild != 0)
				{
					for(int i = 0; i < child.length; i++)
					{
						if(child[i] != null)
						{
							child[i].pruneLevel();
						}
					}
				}
				
				if(level == cube.depth)
				{
					pruneChild();
				}
			}

			/**
             * Remove any nodes that have fewer than threshold
			 * pixels. Also, as long as we're walking the tree:
			 *
			 *  - figure out the color with the fewest pixels
             *  - recalculate the total number of colors in the tree
             */
			long reduce(long threshold, long next_threshold)
			{
				if(nchild != 0)
				{
					for(int i = 0; i < child.length; i++)
					{
						if(child[i] != null)
						{
							next_threshold = child[i].reduce(threshold, next_threshold);
						}
					}
				}
				
				if(number_pixels <= threshold)
				{
					pruneChild();
				}
				else
				{
					if(unique != 0)
					{
						cube.colors++;
					}
					
					if(number_pixels < next_threshold)
					{
						next_threshold = number_pixels;
					}
				}
				
				return next_threshold;
			}

			/*
             * colormap traverses the color cube tree and notes each
             * colormap entry. A colormap entry is any node in the
             * color cube tree where the number of unique colors is
             * not zero.
             */
			void colormap()
			{
				if(nchild != 0)
				{
					for(int i = 0; i < child.length; i++)
					{
						if(child[i] != null)
						{
							child[i].colormap();
						}
					}
				}
				
				if(unique != 0)
				{
					int r = (int)((total_red + (unique >> 1)) / unique);
					int g = (int)((total_green + (unique >> 1)) / unique);
					int b = (int)((total_blue + (unique >> 1)) / unique);
					int a = (int)((total_alpha + (unique >> 1)) / unique);
					
					cube.colormap[cube.colors] = (((a & 0xFF) << 24) |
												  ((r & 0xFF) << 16) |
												  ((g & 0xFF) << 8) |
												  ((b & 0xFF) << 0));
					
					color_number = cube.colors++;
				}
			}

			/* ClosestColor traverses the color cube tree at a
             * particular node and determines which colormap entry
             * best represents the input color.
             */
			void closestColor(int red, int green, int blue, int alpha, Search search)
			{
				if(nchild != 0)
				{
					for(int i = 0; i < child.length; i++)
					{
						if(child[i] != null)
						{
							child[i].closestColor(red, green, blue, alpha, search);
						}
					}
				}

				if(unique != 0)
				{
					int color = cube.colormap[color_number];
					int distance = cube.usealpha ? distance(color, red, green, blue, alpha) : distance(color, red, green, blue);
					
					if(distance < search.distance)
					{
						search.distance = distance;
						search.color_number = color_number;
					}
				}
			}

			/**
			 * Figure out the distance between this node and some color.
			 */
			final static int distance(int color, int r, int g, int b)
			{
				return (SQUARES[((color >> 16) & 0xFF) - r + MAX_RGB] +
						SQUARES[((color >> 8) & 0xFF) - g + MAX_RGB] +
						SQUARES[((color >> 0) & 0xFF) - b + MAX_RGB]);
			}

			/**
			 * Figure out the distance between this node and some color.
			 */
			final static int distance(int color, int r, int g, int b, int a)
			{
				return (SQUARES[((color >> 16) & 0xFF) - r + MAX_RGB] +
						SQUARES[((color >> 8) & 0xFF) - g + MAX_RGB] +
						SQUARES[((color >> 0) & 0xFF) - b + MAX_RGB] +
						SQUARES[((color >> 24) & 0xFF) - a + MAX_RGB]);
			}

			public String toString()
			{
				StringBuffer buf = new StringBuffer();
				
				if(parent == this)
				{
					buf.append("root");
				}
				else
				{
					buf.append("node");
				}
				
				buf.append(' ');
				buf.append(level);
				buf.append(" [");
				buf.append(mid_red);
				buf.append(',');
				buf.append(mid_green);
				buf.append(',');
				buf.append(mid_blue);
				buf.append(']');
				
				return new String(buf);
			}
		}
	}
}
