package ie.dcu.auto.automator;

import static ie.dcu.auto.Constants.ERROR_VALUE;
import ie.dcu.auto.*;
import ie.dcu.graph.dijkstra.IntGraph;
import ie.dcu.image.dt.DistanceTransform;
import ie.dcu.matrix.*;
import ie.dcu.segment.annotate.*;
import ie.dcu.stats.InversionMethod;

import java.util.*;

import org.eclipse.swt.graphics.Point;

/**
 * Implementation of the more complex non-deterministic strategy for interactive
 * segmentation evaluation. This automator weights the pixels so that the paths
 * found tend to stay closer to the center of the shape, rather than using the
 * absolute shortest path
 * 
 * @author Kevin McGuinness
 * 
 */
public class NonDeterministicAutomator3 
	extends AbstractAutomator implements Automator 
{
	private final Set<Index2D> seeds;
	
	public NonDeterministicAutomator3(AutomationData data) {
		super(data);
		this.seeds = new HashSet<Index2D>();
	}

	@Override
	protected void doReset() {
		seeds.clear();
	}

	@Override
	protected void doStart() {
		ByteMatrix fgError = getForegroundError();
		ByteMatrix bgError = getBackgroundError();

		// Distance transform
		DoubleMatrix fgDT = distanceTransform(fgError);
		DoubleMatrix bgDT = distanceTransform(bgError);
		
		// Initialize foreground probability field
		List<Integer> fgPath = createPath(fgError, fgDT);
		List<Integer> bgPath = createPath(bgError, bgDT);
		
		// Add seeds
		if (fgPath != null && bgPath != null) {
			addSeedPath(fgDT, fgPath, AnnotationType.Foreground);
			addSeedPath(bgDT, bgPath, AnnotationType.Background);
		} else if (fgPath != null) {
			addSeedPath(fgDT, fgPath, AnnotationType.Foreground);
		} else if (bgPath != null) {
			addSeedPath(bgDT, bgPath, AnnotationType.Background);
		} else {
			log("Automation complete");
			setDone(true);
		}
	}

	@Override
	protected void doStep() {
		ByteMatrix fgError = getForegroundError();
		ByteMatrix bgError = getBackgroundError();
		
		DoubleMatrix fgDT = distanceTransform(fgError);
		DoubleMatrix bgDT = distanceTransform(bgError);
		
		DoubleMatrix field = getProbabilityField(fgDT, bgDT);
		
		if (field != null) {
			InversionMethod rv = new InversionMethod(field.values);
			int seed = rv.random();
			List<Integer> path = null;
			AnnotationType type;
			
			if (fgError.values[seed] == ERROR_VALUE) {
				type = AnnotationType.Foreground;
				path = createPath(fgError, fgDT, seed);
			
			} else {
				type = AnnotationType.Background;
				path = createPath(bgError, bgDT, seed);	
			}
			
			if (path != null) {
				addSeedPath(field, path, type);
			} else {
				addSeed(field, seed, type);
			}
			
		} else {
			log("Automation complete");
			setDone(true);
		}
	}

	private void addSeed(DoubleMatrix field, int seed, AnnotationType type) {
		log("Adding single seed");
		
		Index2D index = field.indexOf(seed);
		double distance = field.doubleAt(seed);
		int size = (int) Math.min(distance/2, MAX_BRUSH_SIZE);
		size = Math.max(size, 1);
		Point pt = getSwtPointForIndex(index);
		getAnnotations().add(new Annotation(type, size, pt));
		seeds.add(index);
	}

	private void addSeedPath(
		DoubleMatrix field, List<Integer> path, AnnotationType type) {
		
		log("Adding seed path");
		
		int size = MAX_BRUSH_SIZE;
		for (int k : path) {
			double distance = field.doubleAt(k);
			size = (int) Math.min(size, distance/2);
		}
		
		size = Math.max(size, 1);
		
		for (int k : path) {
			Index2D index = field.indexOf(k);
			Point pt = getSwtPointForIndex(index);
			getAnnotations().add(new Annotation(type, size, pt));
			seeds.add(index);
		}
	}

	private List<Integer> createPath(ByteMatrix error, DoubleMatrix dt) {
		if (nonzero(dt)) {
			
			// Select random variable
			InversionMethod field = new InversionMethod(dt.values);
			int source = field.random();
			
			// Create path
			return createPath(error, dt, source);
		}
		
		return null;
	}

	private List<Integer> createPath(ByteMatrix error, DoubleMatrix dt, int source) {
		
		// Build graph
		IntGraph graph = buildGraph(error, dt);
		
		// The graph may not contain the source pixel if it has no neighbors
		if (graph.hasNode(source)) {
			
			// Ok, it does
			graph.findPaths(source);
			
			// Mark inaccessible pixels with probability zero
			for (int i = 0; i < dt.size; i++) {
				if (!graph.hasNode(i)) {
					dt.values[i] = 0.0;
				}
			}
			
			if (nonzero(dt)) {
				
				// Select random variable
				InversionMethod field = new InversionMethod(dt.values);
				int target = field.random();
				
				// Get shortest path
				return graph.getPathFrom(target);
			}
		} 
		
		return null;
	}
	
	private DoubleMatrix getProbabilityField(
		DoubleMatrix fgDT, DoubleMatrix bgDT) {
		
		// Sum transforms
		DoubleMatrix field = fgDT.clone();
		for (int i = 0; i < field.size; i++) {
			field.values[i] += bgDT.values[i];
		}
		
		// Set probability of selected seeds to zero
		for (Index2D seed : seeds) {
			field.setDoubleAt(seed, 0.0);
		}
		
		if (!isStoppingCriteriaSatisfied(field)) {
			return field;
		}
		
		return null;
	}
	
	private static IntGraph buildGraph(ByteMatrix matrix, DoubleMatrix dt) {
		
		final double dtmax = dt.maxValue();
		final double SQRT_2 = Math.sqrt(2.0);
		final byte[] values = matrix.values;
		final int rows = matrix.rows;
		final int cols = matrix.cols;
		final IntGraph graph = new IntGraph(matrix.size);
		
		for (int i = 0; i < rows; i++) {
			int k = i * cols;

			for (int j = 0; j < cols; j++, k++) {
				if (values[k] != ERROR_VALUE) {
					continue;
				}
				
				// Locations neighboring pixels
				final int north = k - cols;
				final int south = k + cols;
				final int west  = k - 1;
				final int east  = k + 1;
				final int northWest = north - 1;
				final int northEast = north + 1;
				final int southWest = south - 1;
				final int southEast = south + 1;
				
				final boolean hasNorth = i > 0;
				final boolean hasSouth = i + 1 < rows;
				final boolean hasWest  = j > 0;
				final boolean hasEast  = j + 1 < cols;
				
				double w;
				
				if (hasNorth) {
					if (hasWest && values[northWest] == ERROR_VALUE) {
						w = weight(dt.values[northWest], dtmax, SQRT_2);
						graph.addDirectedEdge(k, northWest, w);
					}
					
					if (values[north] == ERROR_VALUE) {
						w = weight(dt.values[north], dtmax, 1);
						graph.addDirectedEdge(k, north, w);
					}
				
					if (hasEast && values[northEast] == ERROR_VALUE) {
						w = weight(dt.values[northEast], dtmax, SQRT_2);
						graph.addDirectedEdge(k, northEast, w);
					}
				}
				
				if (hasWest && values[west] == ERROR_VALUE) {
					w = weight(dt.values[west], dtmax, 1);
					graph.addDirectedEdge(k, west, w);
				}
				
				if (hasEast && values[east] == ERROR_VALUE) {
					w = weight(dt.values[east], dtmax, 1);
					graph.addDirectedEdge(k, east, w);
				}
				
				if (hasSouth) {
					if (hasWest && values[southWest] == ERROR_VALUE) {
						w = weight(dt.values[southWest], dtmax, SQRT_2);
						graph.addDirectedEdge(k, southWest, w);
					}
					
					if (values[south] == ERROR_VALUE) {
						w = weight(dt.values[south], dtmax, 1);
						graph.addDirectedEdge(k, south, w);
					}
					
					if (hasEast && values[southEast] == ERROR_VALUE) {
						w = weight(dt.values[southEast], dtmax, SQRT_2);
						graph.addDirectedEdge(k, southEast, w);
					}
				}
			}
		}
		
		return graph;
	}
	
	private static double weight(double n, double m, double d) {
		return d*Math.exp(-n/m);
	}

	private static DoubleMatrix distanceTransform(ByteMatrix error) {
		DistanceTransform dt = new DistanceTransform();
		dt.init(error, Constants.ERROR_VALUE);
		return dt.computeTransform();
	}
	
	private static boolean isStoppingCriteriaSatisfied(DoubleMatrix dt) {
		return dt.maxValue() < 2.0;
	}
	
	private static boolean nonzero(DoubleMatrix matrix) {
		for (int i = 0; i < matrix.size; i++) {
			if (matrix.values[i] != 0) {
				return true;
			}
		}
		return false;
	}
	
	public static boolean isNonDeterministic() {
		return true;
	}
}