package helpers;

import java.awt.AlphaComposite;
import java.awt.Color;
import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.awt.image.RenderedImage;
import java.io.File;
import java.io.IOException;
import java.util.Set;

import javax.imageio.ImageIO;
import javax.media.jai.iterator.RandomIter;
import javax.media.jai.iterator.RandomIterFactory;

import models.Image;
import play.Logger;

public class ImageHelper {

	private static final int MAX_HEIGHT = 600;
	private static final int MAX_WIDTH = 600;
	private static final int MAX_SOBEL_HEIGHT = 1024;
	private static final int MAX_SOBEL_WIDTH = 1024;
	private static final int REF_DIM = 300;

	private static final int COLOR_COUNT_LIMIT = 100;

	/**
	 * This creates a thumbnail image with max 600px w/h.
	 * 
	 * @param image
	 *            the image object
	 * @return the thumbnail file.
	 */
	public static File thumbnail(final Image image) {
		File imageFile = new File(image.imagePath);
		File thumbnail = new File(image.imagePath
				.replace("." + image.extension, "_small.png")
				.replace(FileHelper.FILE_DIRECTORY, FileHelper.THUMBS_DIRECTORY));
		try {
			BufferedImage bimg = ImageIO.read(imageFile);
			thumbnail = resizeImage(bimg, thumbnail, MAX_WIDTH, MAX_HEIGHT);
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}
		return thumbnail;
		
	}
	
	/**
	 * This creates a normalized reference Image with REF_DIM*REF_DIM
	 * dimensions.
	 * 
	 * @return
	 */
	public static File referenceImage(final Image image) {
		File imageFile = new File(image.imagePath);
		File refImage = new File(image.imagePath.replace("." + image.extension, "_ref.png")
				.replace(FileHelper.FILE_DIRECTORY, FileHelper.REFS_DIRECTORY));
		try {
			BufferedImage bimg = ImageIO.read(imageFile);
			Logger.info("sizing {} from {}/{} to {}/{}", imageFile.getAbsolutePath(), bimg.getWidth(), bimg.getHeight(),
					REF_DIM, REF_DIM);
			BufferedImage scaledBI = new BufferedImage(REF_DIM, REF_DIM, BufferedImage.TYPE_INT_RGB);
			Graphics2D g = scaledBI.createGraphics();
			g.setComposite(AlphaComposite.Src);
			g.drawImage(bimg, 0, 0, REF_DIM, REF_DIM, null);
			g.dispose();
			ImageIO.write(scaledBI, "png", refImage);
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}
		return refImage;
	}

	public static File edgeImage(final Image image) {
		File imageFile = new File(image.imagePath);
		File sobelImage = new File(image.imagePath.replace("." + image.extension, "_sobel.png")
				.replace(FileHelper.FILE_DIRECTORY, FileHelper.SOBEL_DIRECTORY));
		return createSobel(imageFile, sobelImage);
	}
	
	public static void countColors(final Image image) {
		File file = new File(image.thumbnailPath);
		try {
			BufferedImage bimg = ImageIO.read(file);
			CountMap cm = getColorsCount(bimg);
			addColors(image, cm);
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}
	}

	public static void setImageColors(final Image img, CountMap countMap) {
		Image image = Image.Find.where().eq("checksum", img.checksum).findUnique();
		models.ImageColor ic = image != null && image.colors.size() > 0 ? image.colors.get(0) : null;
		if (ic != null) {
			image.mainColor = ic.color.toString();
			Logger.info("image {}: main color is {}", image.id, image.mainColor);
			models.Color avgColor = models.Color.findOrCreate(countMap.getAvgR(), countMap.getAvgG(), countMap.getAvgB());
			image.avgColor = avgColor.toString();
			Logger.info("image {}: avg color is {}", image.id, image.avgColor);
			int medIndex = getMedianColorKey(countMap.getCounts().keySet(), countMap.getSize() / 2);
			Logger.info("image {}: median index is {}", image.id, medIndex);
			Color realMedColor = new Color(countMap.getCounts().get(medIndex), true);
			models.Color medColor = models.Color.findOrCreate(realMedColor.getRed(), realMedColor.getGreen(), realMedColor.getBlue());
			image.medColor = medColor.toString();
			Logger.info("image {}: med color is {}", image.id, image.medColor);
			image.save();
		}
	}
	
	public static Color[][] createSignature(final Image img) {
		try {
			File reference = new File(img.refImagePath);
			if(reference.exists()) {
				RenderedImage ref = ImageIO.read(reference);
				Color[][] sig = calcSignature(ref);
				Logger.info("image {}: sig {}", img.id, sig);
				return sig;
			}
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}	
		return null;
	}

	private static Integer getMedianColorKey(Set<Integer> keys, int index) {
		int count = 0;
		for (Integer k : keys) {
			count++;
			if (count == index) {
				return k;
			}
		}
		return null;
	}

	private static CountMap getColorsCount(BufferedImage image) {
		CountMap cm = new CountMap();
		int w = image.getWidth();
		int h = image.getHeight();
		for (int y = 0; y < h; y++) {
			for (int x = 0; x < w; x++) {
				cm.add(image.getRGB(x, y));
			}
		}
		return cm;
	}

	private static void addColors(final Image image, CountMap countMap) {
		Logger.info("image {}: adding {} colors", image.id, countMap.getSize());
		for (Integer rgb : countMap.getCounts().keySet()) {
			if (countMap.getCounts().get(rgb) > COLOR_COUNT_LIMIT) {
				Color rgbColor = new Color(rgb, true);
				models.Color c = models.Color.findOrCreate(rgbColor.getRed(), rgbColor.getGreen(), rgbColor.getBlue());
				models.ImageColor ic = models.ImageColor.findOrCreate(image, c);
				ic.count = countMap.getCounts().get(rgb);
				ic.save();
			}
		}
		Logger.info("image {}: done!", image.id);
		Logger.info("image {}: stats {} ColorCounts", image.id, countMap.getSize());
		ImageHelper.setImageColors(image, countMap);
	}
	
	/**
	 * Creates an image with sobel operation based edge detection.
	 * 
	 * @see https://stackoverflow.com/a/30511674
	 */
	private static File createSobel(File inputFile, File outputFile) {
		try {
			int[][] pixelMatrix = new int[3][3];
			BufferedImage img = ImageIO.read(inputFile);
			BufferedImage outputImg = new BufferedImage(img.getWidth(), img.getHeight(), BufferedImage.TYPE_INT_RGB);
			for (int i = 1; i < img.getWidth() - 1; i++) {
				for (int j = 1; j < img.getHeight() - 1; j++) {
					pixelMatrix[0][0] = new Color(img.getRGB(i - 1, j - 1)).getRed();
					pixelMatrix[0][1] = new Color(img.getRGB(i - 1, j)).getRed();
					pixelMatrix[0][2] = new Color(img.getRGB(i - 1, j + 1)).getRed();
					pixelMatrix[1][0] = new Color(img.getRGB(i, j - 1)).getRed();
					pixelMatrix[1][2] = new Color(img.getRGB(i, j + 1)).getRed();
					pixelMatrix[2][0] = new Color(img.getRGB(i + 1, j - 1)).getRed();
					pixelMatrix[2][1] = new Color(img.getRGB(i + 1, j)).getRed();
					pixelMatrix[2][2] = new Color(img.getRGB(i + 1, j + 1)).getRed();

					int edge = (int) convolution(pixelMatrix);
					outputImg.setRGB(i, j, (edge << 16 | edge << 8 | edge));
				}
			}
			outputFile = resizeImage(outputImg, outputFile, MAX_SOBEL_WIDTH, MAX_SOBEL_HEIGHT );
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}
		return outputFile;
	}
	
	private static double convolution(int[][] pixelMatrix) {
		int gy = (pixelMatrix[0][0] * -1) + (pixelMatrix[0][1] * -2) + (pixelMatrix[0][2] * -1) + (pixelMatrix[2][0])
				+ (pixelMatrix[2][1] * 2) + (pixelMatrix[2][2] * 1);
		int gx = (pixelMatrix[0][0]) + (pixelMatrix[0][2] * -1) + (pixelMatrix[1][0] * 2) + (pixelMatrix[1][2] * -2)
				+ (pixelMatrix[2][0]) + (pixelMatrix[2][2] * -1);
		return Math.sqrt(Math.pow(gy, 2) + Math.pow(gx, 2));

	}
	
	/**
	 * Calculates the signature of an image.
	 * 
	 * @see http://150.163.105.214:8080/JIPCookbook/6050-howto-compareimages.jsp#howtohowdoicomparetwoimagestoseeiftheyareequal
	 */
	private static Color[][] calcSignature(RenderedImage img) {
		// Get memory for the signature.
		Color[][] sig = new Color[5][5];
		// For each of the 25 signature values average the pixels around it.
		// Note that the coordinate of the central pixel is in proportions.
		float[] prop = new float[] { 1f / 10f, 3f / 10f, 5f / 10f, 7f / 10f, 9f / 10f };
		for (int x = 0; x < 5; x++) {
			for (int y = 0; y < 5; y++) {
				sig[x][y] = averageAround(img, prop[x], prop[y]);
			}
		}
		return sig;
	}
	
	
	/**
	 * This method calculates the distance between the signatures of an image
	 * and the reference one. The signatures for the image passed as the
	 * parameter are calculated inside the method.
	 * 
	 * @see http://150.163.105.214:8080/JIPCookbook/6050-howto-compareimages.jsp#howtohowdoicomparetwoimagestoseeiftheyareequal
	 */
	private static double calcDistance(RenderedImage otherImg, Color[][] signature) {
		// Calculate the signature for that image.
		Color[][] sigOther = calcSignature(otherImg);
		// There are several ways to calculate distances between two vectors,
		// we will calculate the sum of the distances between the RGB values of
		// pixels in the same positions.
		double dist = 0;
		for (int x = 0; x < 5; x++)
			for (int y = 0; y < 5; y++) {
				int r1 = signature[x][y].getRed();
				int g1 = signature[x][y].getGreen();
				int b1 = signature[x][y].getBlue();
				int r2 = sigOther[x][y].getRed();
				int g2 = sigOther[x][y].getGreen();
				int b2 = sigOther[x][y].getBlue();
				double tempDist = Math.sqrt((r1 - r2) * (r1 - r2) + (g1 - g2) * (g1 - g2) + (b1 - b2) * (b1 - b2));
				dist += tempDist;
			}
		return dist;
	}
	    
	/**
	 * This method averages the pixel values around a central point and return the average as an instance of Color. 
	 * The point coordinates are proportional to the image.
	 * 
	 * @see http://150.163.105.214:8080/JIPCookbook/6050-howto-compareimages.jsp#howtohowdoicomparetwoimagestoseeiftheyareequal
	 */
	private static Color averageAround(RenderedImage img, double px, double py) {

		// Get an iterator for the image.
		RandomIter iterator = RandomIterFactory.create(img, null);
		// Get memory for a pixel and for the accumulator.
		double[] pixel = new double[3];
		double[] accum = new double[3];
		// The size of the sampling area.
		int sampleSize = 15;
		int numPixels = 0;
		// Sample the pixels.
		for (double x = px * REF_DIM - sampleSize; x < px * REF_DIM + sampleSize; x++) {
			for (double y = py * REF_DIM - sampleSize; y < py * REF_DIM + sampleSize; y++) {
				iterator.getPixel((int) x, (int) y, pixel);
				accum[0] += pixel[0];
				accum[1] += pixel[1];
				accum[2] += pixel[2];
				numPixels++;
			}
		}
		// Average the accumulated values.
		accum[0] /= numPixels;
		accum[1] /= numPixels;
		accum[2] /= numPixels;
		return new Color((int) accum[0], (int) accum[1], (int) accum[2]);
	}
	
	private static File resizeImage(BufferedImage bimg, File output, int maxWidth, int maxHeight) {
		int originalH, originalW, scaledH, scaledW;
		try {
			originalH = bimg.getHeight();
			originalW = bimg.getWidth();

			if (originalH > originalW) {
				scaledH = maxHeight;
				scaledW = (int) (maxWidth * (float) originalW / (float) originalH);
			} else {
				scaledW = maxWidth;
				scaledH = (int) (maxHeight * (float) originalH / (float) originalW);
			}
			Logger.info("sizing {} from {}/{} to {}/{}", output.getAbsolutePath(), originalW, originalH, scaledW, scaledH);

			BufferedImage scaledBI = new BufferedImage(scaledW, scaledH, BufferedImage.TYPE_INT_RGB);
			Graphics2D g = scaledBI.createGraphics();
			g.setComposite(AlphaComposite.Src);
			g.drawImage(bimg, 0, 0, scaledW, scaledH, null);
			g.dispose();
			ImageIO.write(scaledBI, "png", output);
		} catch (IOException e) {
			Logger.error(e.getLocalizedMessage());
		}
		return output;		
	}	
}