/*
 *  Brno University of Technology
 *
 *  Copyright (C) 2009-2010 by Brno University of Technology and the contributors
 *
 *  Complete list of developers available at our web site:
 *
 *       http://spl.utko.feec.vutbr.cz
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see http://www.gnu.org/licenses/.
 */
package cz.vutbr.feec.imageprocessing.op.segmentation.watershed;

import ij.ImagePlus;
import ij.gui.ImageWindow;
import ij.process.ByteProcessor;
import ij.process.ColorProcessor;
import ij.process.ImageProcessor;

import java.awt.Color;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;

import com.rapidminer.parameter.UndefinedParameterError;

import cz.vutbr.feec.imageprocessing.IOObject.ImagePlusSegmentedIOObject;
import cz.vutbr.feec.imageprocessing.IOObject.innerObject.SegmentMask;
import cz.vutbr.feec.imageprocessing.op.segmentation.watershed.WPoint.LblType;

/**
 * The Class WatershedBase
 * 
 * @param <E>
 * @MISC{Roerdink00thewatershed, author = {Jos B. T. M. Roerdink and Arnold
 *                               Meijster}, title = {The Watershed Transform:
 *                               Definitions, Algorithms and Parallelization
 *                               Strategies}, year = {2000} }
 */
public class WatershedAlgorithm {
	private ImagePlus tempImage = null;
	private WPoint curp;
	private int curdist = 0;
	private int curlab = 0;
	private int width;
	private int height;
	private ImagePlus inputImage = null;
	private ImageStructure imStr;
	private Region[] regions;
	private boolean debug;

	
	public ImagePlus doWork(ImagePlus inputImage, RgbToGSModel convertModel,
			ConnectivityModels connectivityModel, int smallRegionThr,
			int mergingThr, boolean wsheds) {
		this.inputImage =inputImage;
		height = inputImage.getHeight();
		width = inputImage.getWidth();

		// copy pixels to data structure
		imStr = new ImageStructure(inputImage, convertModel, connectivityModel);
		//in debug mode show supporting informations 
		if (debug) showProgress(imStr, "Original Image", true, false, false);

		//main part - segmentation of an image
		segmentate();
		
		//merging regions to met needed results
		postprocessing(wsheds, smallRegionThr, mergingThr);
		
		//deliver segmented image as ImagePlus
		return imStr.deliverResultImage(regions, false);
	}

	/**
	 * Vincent-Soille fast watershed segmentation algorithm @
	 */
	private void segmentate() {
		LinkedList<WPoint> fifo = new LinkedList<WPoint>();
		boolean possibleWshed = false;

		//for debug
		// showProgress(imStr, "input image");

		// start flooding
		for (int h = 0; h < 256; h++) {

			for (int x = 0; x < width; x++) {
				for (int y = 0; y < height; y++) {
					curp = imStr.getPoint(x, y);
					if (curp.getGray() == h) {
						curp.setLabelToMask();
						// for debug
						// curp.setColor(255, 0, 255);
						if (curp.hasLabeledNeighbor()) {
							curp.setDistance(1);
							fifo.add(curp);
						}
					}
				}
			}
			curdist = 1;
			fifo.add(new WPoint());
			//for debug
			// showProgress(imStr, String.format("H:%d", h));
			while (!fifo.isEmpty()) {
				curp = fifo.poll();
				if (curp.getLabel() == LblType.FICTIVE) {
					if (fifo.isEmpty())
						break;
					else {
						fifo.add(new WPoint());
						curdist++;
						curp = fifo.poll();
					}
				}

				//for debug
				// curp.setColor(0,255,255);
				// showProgress(imStr, String.format("H:%d,L:%d", h,curlab));

				for (WPoint neighbor : curp.getNeighbors()) {
					if (neighbor.getDistance() < curdist
							&& (neighbor.isLabelOverZero() || neighbor.isWShed())) {
						if (neighbor.isLabelOverZero()) {
							if (curp.isMask()) {
								curp.setLabel(neighbor.getLabel());
								curp.setDistance(0);
								//for debug
								// curp.setColor(255, 255, 0);
								// showProgress(imStr,
								// String.format("H:%d,L:%d", h,curlab));
							} else if (curp.getLabel() != neighbor.getLabel()) {
								curp.setLabelToWshed();
								//for debug
								// curp.setColor(0, 255, 0);
								// showProgress(imStr,
								// String.format("H:%d,L:%d", h,curlab));
							}
						} else if (curp.isMask()) {
							possibleWshed = true;
						}
					} else if (neighbor.isMask() && neighbor.getDistance() == 0) {
						neighbor.setDistance(curdist + 1);
						fifo.add(neighbor);
					}
				}

				if (possibleWshed && curp.getLabel() == LblType.MASK) {
					curp.setLabelToWshed();
					//for debug
					// curp.setColor(0, 255, 0);
					// showProgress(imStr, String.format("H:%d,L:%d",
					// h,curlab));
				}

			}

			for (int x = 0; x < width; x++) {
				for (int y = 0; y < height; y++) {
					curp = imStr.getPoint(x, y);
					 if (curp.getGray() == h) {
						curp.setDistance(0);
						if (curp.getLabel() == LblType.MASK) {
							curlab++;
							fifo.add(curp);
							curp.setLabel(curlab);
							/*
							 * for debug
							 * curp.setColor(0, 0, 255);
							 * showProgress(imStr, String.format("H:%d,L:%d", h,curlab));
							 */
							while (!fifo.isEmpty()) {
								WPoint curp = fifo.poll();
								ArrayList<WPoint> neighborList = curp
										.getNeighbors();
								for (int n = 0; n < neighborList.size(); n++) {
									WPoint neighbor = neighborList.get(n);
									if (neighbor.getLabel() == LblType.MASK) {
										fifo.add(neighbor);
										neighbor.setLabel(curlab);
										/*
										 *for debug
										 *neighbor.setColor(0, 0, 255);
										 *showProgress(imStr, String.format("H:%d,L:%d", h,curlab));  
										*/
									}
								}
							}
							//for debug
							// showProgress(imStr, String.format("H:%d,L:%d",
							// h,curlab));
						}
					}
				}
			}
			/*
			 * for debug
			 * if (h == 255) showProgress(imStr, String.format("H:%d,L:%d", h,curlab));
			*/

		}
	}

	//for testing
	private void showProgress(ImageStructure imStr, String progressMessage) {
		showProgress(imStr, progressMessage, false, false, false);
	}

	//for testing
	private void showProgress(ImageStructure imStr, String progressMessage,
			Boolean doNotClose, Boolean showRegions, Boolean showWsheds) {
		int left = -1680;
		if (tempImage != null) {
			if (!doNotClose) {
				tempImage.close();
			} else {
				left = tempImage.getWindow().getX() + tempImage.getWidth() + 20;
			}
		}

		if (showRegions) {
			tempImage = imStr.deliverResultImage(regions, showWsheds);
		} else {
			tempImage = imStr.deliverResultImage(null, showWsheds);
		}
		tempImage.setProcessor(tempImage.getProcessor().resize(400));
		tempImage.setTitle(progressMessage);
		tempImage.show(progressMessage);
		ImageWindow w = tempImage.getWindow();
		w.setLocation(left, 400);
		w.setAlwaysOnTop(true);
	}

	private void postprocessing(boolean wsheds, int smallRegionThr,
			int mergingThr) {
		// creates structure for regions according to number of labels
		regions = new Region[curlab + 1];
		for (int i = 0; i <= curlab; i++) {
			regions[i] = new Region(i);
		}

		// fifo1 is used for assigning wsheds to neighboring areas
		LinkedList<WPoint> fifo1 = new LinkedList<WPoint>();
		// fifo2 is used for merging regions
		LinkedList<WPoint> fifo2 = new LinkedList<WPoint>();
		for (int x = 0; x < width; x++) {
			for (int y = 0; y < height; y++) {
				curp = imStr.getPoint(x, y);
				if (curp.isWShed()) {
					fifo1.add(curp);
					fifo2.add(curp);

				} else if (curp.isLabelOverZero())
					regions[curp.getLabel()].addPoint(curp);
			}
		}

		assignWatersheds(fifo1);
		if (debug) {
			showProgress(imStr, regions.length
					+ " regions in total. Before connecting small regions.",
					true, true, true);
		}
		int connectedSmall = connectSmallRegions(smallRegionThr);
		if (debug) {
			showProgress(imStr, (regions.length - connectedSmall)
					+ " regions after connecting small regions.", true, true,
					false);
		}
		int merged = mergeNeighboringRegions(mergingThr, fifo2);
		// imStr.deliverResultImage(null, false).show();
		if (debug) {
			showProgress(imStr, (regions.length - connectedSmall - merged)
					+ " regions after merging. Final result.", true, true,
					false);
		}
		// imStr.deliverResultImage(null, false).show();
	}

	private void setNeighboringRegions(ArrayList<Region> regionsAround) {
		int size = regionsAround.size();
		for (int i = 0; i < size; i++) {
			Region curReg = regionsAround.get(i);
			regionsAround.remove(i);
			for (Region region : regionsAround) {
				region.addNeigboringRegion(curReg);
			}
			regionsAround.add(i, curReg);
		}
	}

	private void assignWatersheds(LinkedList<WPoint> fifo1) {
		while (!fifo1.isEmpty()) {
			WPoint curp = fifo1.poll();
			ArrayList<Region> regionsAround = curp.findRegionsAround(regions);
			setNeighboringRegions(regionsAround);
			search: {
				for (WPoint neighbor : curp.getNeighbors()) {
					if (neighbor.isLabelOverZero() && !neighbor.wasWaterShad()) {
						regions[neighbor.getLabel()].addPoint(curp);
						break search;
					}
				}
				// solve leftovers
				for (WPoint neighbor : curp.getNeighbors()) {
					if (neighbor.isLabelOverZero() && !neighbor.isWShed()) {
						regions[neighbor.getLabel()].addPoint(curp);
						break search;
					}
				}
			}
		}
	}

	private int connectSmallRegions(int smallRegionThr) {
		int numConnected = 0;
		int small = (int) width * height / 1000 * smallRegionThr;

		for (Region smallReg : regions) {
			if (smallReg.getNumOfPointsInRegion() < small
					&& smallReg.getNumOfPointsInRegion() > 0) {
				Region similarRegion = smallReg.findMostSimilarRegion();
				merge2Regions(similarRegion, smallReg);
				numConnected++;
			}
		}
		return numConnected;
	}

	private int mergeNeighboringRegions(int mergingThr, LinkedList<WPoint> fifo2) {
		int numConnected = 0;
		for (Region region : regions) {
			if (region.getNumOfPointsInRegion() > 0) {

				ArrayList<Region> neighRegs = region.getNeighboringRegions();
				for (int i = 0; i < neighRegs.size();) {
					Region neigRegion = neighRegs.get(i);
					if (region.isRegionSimilar(mergingThr, neigRegion)) {
						merge2Regions(region, neigRegion);
						numConnected++;
					} else {
						i++;
					}
				}
			}
		}
		return numConnected;
	}

	private static void merge2Regions(Region region1, Region region2) {
		region1.connect(region2);
		region1.remapNeighboringRegions(region2);
	}

	public void debug(boolean b) {
		this.debug = b;

	}

	public ImagePlusSegmentedIOObject setOutput(ImagePlusSegmentedIOObject segmentedImg) throws UndefinedParameterError {
			
			int iWidth = segmentedImg.getImage().getWidth();
			int iHeight = segmentedImg.getImage().getHeight();
			int i = 0;
			
			for (Region region : regions ){
				int size = region.getNumOfPointsInRegion();
				if (size <= 0)
					continue;
				
				// create new mask for each region by default all pixels are black
				ImageProcessor maskedImageProc = new ColorProcessor(iWidth, iHeight);
				maskedImageProc.setColor(Integer.MAX_VALUE);
				//set white for each pixel within the region 
				ArrayList<WPoint> points = region.getPoints();
				for (WPoint point:points) {
					int x = point.getPosX();
					int y = point.getPosY();
					maskedImageProc.drawDot(x, y);
				}
				
				ImagePlus maskedImage = new ImagePlus("mask " + i, maskedImageProc);
				
				// add new mask and set relationships
				SegmentMask sm = new SegmentMask(i, maskedImage);
				sm.setSize(size);
				sm.setLevel(0);
				segmentedImg.addSegmentMask(sm);
				i++;
			}
			
			return segmentedImg;
	}

}
