package ie.dcu.auto.automator;

import ie.dcu.auto.*;
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;

/**
 * Automates the evaluation of interactive segmentation algorithms using 
 * a simple non-deterministic method.
 * 
 * @author Kevin McGuinness
 */
public class NonDeterministicAutomator1 
	extends AbstractAutomator implements Automator 
{
	private final Set<Index2D> seeds;

	public NonDeterministicAutomator1(AutomationData data) {
		super(data);
		this.seeds = new HashSet<Index2D>();
	}
	
	protected void doReset() {
		seeds.clear();
	}
	
	protected void doStart() {
		
		// Distance transform
		DoubleMatrix fgDT = distanceTransform(getForegroundError());
		DoubleMatrix bgDT = distanceTransform(getBackgroundError());
		
		// Initialize foreground probability field
		Index2D fgIndex = null;
		if (nonzero(fgDT)) {
			InversionMethod fgProb = new InversionMethod(fgDT.values);
			fgIndex = fgDT.indexOf(fgProb.random());
		}
		
		// Initialize background probability field
		Index2D bgIndex = null;
		if (nonzero(bgDT)) {
			InversionMethod bgProb = new InversionMethod(bgDT.values);
			bgIndex = bgDT.indexOf(bgProb.random());
		}
		
		// Add seeds
		if (fgIndex != null && bgIndex != null) {
			addSeed(fgDT, fgIndex, AnnotationType.Foreground);
			addSeed(bgDT, bgIndex, AnnotationType.Background);
		} else if (fgIndex != null) {
			addSeed(fgDT, fgIndex, AnnotationType.Foreground);
		} else if (bgIndex != null) {
			addSeed(bgDT, bgIndex, AnnotationType.Background);
		} else {
			log("Automation complete");
			setDone(true);
		}
	}

	protected void doStep() {
		ByteMatrix fgError = getForegroundError();
		ByteMatrix bgError = getBackgroundError();
		
		DoubleMatrix field = getProbabilityField(fgError, bgError);
		if (field != null) {
			InversionMethod rv = new InversionMethod(field.values);
			Index2D seed = field.indexOf(rv.random());
			AnnotationType type = 
				(fgError.byteAt(seed) == Constants.ERROR_VALUE) ?
				AnnotationType.Foreground : AnnotationType.Background;
			addSeed(field, seed, type);
			
		} else {
			log("Automation complete");
			setDone(true);
		}
	}
	
	private DoubleMatrix getProbabilityField(ByteMatrix fgError, ByteMatrix bgError) {
		
		// Distance transforms
		DoubleMatrix fgDT = distanceTransform(fgError);
		DoubleMatrix bgDT = distanceTransform(bgError);
		
		// Sum transforms
		DoubleMatrix field = fgDT;
		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 void addSeed(DoubleMatrix dt, Index2D seed, AnnotationType type) {
		double distance = dt.doubleAt(seed);
		int size = (int) Math.min(distance/2, MAX_BRUSH_SIZE);
		size = Math.max(size, 1);
		Point pt = getSwtPointForIndex(seed);
		getAnnotations().add(new Annotation(type, size, pt));
		seeds.add(seed);
	}

	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;
	}
}