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 java.util.*;

import org.eclipse.swt.graphics.Point;

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

	public DeterministicAutomator(AutomationData data) {
		super(data);
		this.seeds = new HashSet<Index2D>();
	}
	
	protected void doReset() {
		seeds.clear();
	}
	
	protected void doStart() {
		
		// Get error
		ByteMatrix fgError = getForegroundError();
		ByteMatrix bgError = getBackgroundError();
		
		// DT
		IntMatrix fgDT = distanceTransform(fgError);
		IntMatrix bgDT = distanceTransform(bgError);
		
		// Maxima
		MaximaLocator fgLocator = new MaximaLocator(fgDT);
		List<Index2D> fgMaxima = fgLocator.findAll(2);
		
		MaximaLocator bgLocator = new MaximaLocator(bgDT);
		List<Index2D> bgMaxima = bgLocator.findAll(2);
		
		// Add seeds
		if (fgMaxima.size() > 0 && bgMaxima.size() > 0) {
			addForegroundSeeds(fgMaxima, fgLocator.getMaximum());
			addBackgroundSeeds(bgMaxima, bgLocator.getMaximum());
		} else if (fgMaxima.size() > 0) {
			addForegroundSeeds(fgMaxima, fgLocator.getMaximum());
		} else if (bgMaxima.size() > 0) {
			addBackgroundSeeds(bgMaxima, bgLocator.getMaximum());
		} else {
			log("Automation complete");
			setDone(true);
		}
	}
	
	protected void doStep() {
		
		// Get error
		ByteMatrix fgError = getForegroundError();
		ByteMatrix bgError = getBackgroundError();
		
		// DT
		IntMatrix fgDT = distanceTransform(fgError);
		IntMatrix bgDT = distanceTransform(bgError);
		
		// Maxima
		MaximaLocator fgLocator = new MaximaLocator(fgDT);
		List<Index2D> fgMaxima = fgLocator.findAll(2);
		
		MaximaLocator bgLocator = new MaximaLocator(bgDT);
		List<Index2D> bgMaxima = bgLocator.findAll(2);
		
		log("Found %d potential foreground seeds", fgMaxima.size());
		log("Found %d potential background seeds", bgMaxima.size());
		
		// Remove previously selected seeds
		removeSelectedSeeds(fgMaxima);
		removeSelectedSeeds(bgMaxima);
		
		// Add seeds
		if (fgMaxima.size() > 0 && bgMaxima.size() > 0) {
			
			if (fgLocator.getMaximum() >= bgLocator.getMaximum()) {
				addForegroundSeeds(fgMaxima, fgLocator.getMaximum());	
			} else {
				addBackgroundSeeds(bgMaxima, bgLocator.getMaximum());
			}
			
		} else if (fgMaxima.size() > 0) {		
			addForegroundSeeds(fgMaxima, fgLocator.getMaximum());
			
		} else if (bgMaxima.size() > 0) {
			addBackgroundSeeds(bgMaxima, bgLocator.getMaximum());
			
		} else {
			log("Automation complete");
			setDone(true);
		}
	}

	private void removeSelectedSeeds(List<Index2D> seeds) {
		Iterator<Index2D> iterator = seeds.iterator();
		
		while (iterator.hasNext()) {
			Index2D seed = iterator.next();
			if (this.seeds.contains(seed)) {
				iterator.remove();
			} 
		}
	}

	private void addBackgroundSeeds(List<Index2D> seeds, int maximum) {
		double distance = Math.sqrt(maximum);
		
		log("Adding %d background seeds (distance: %f)", seeds.size(), distance);
		
		int size = (int) Math.min(distance/2, MAX_BRUSH_SIZE);
		
		AnnotationManager annotations = getAnnotations();
		for (Index2D seed : seeds) {
			this.seeds.add(seed);
			Point pt = getSwtPointForIndex(seed);
			annotations.add(new Annotation(AnnotationType.Background, size, pt));
		}
	}

	private void addForegroundSeeds(List<Index2D> seeds, int maximum) {
		double distance = Math.sqrt(maximum);
		
		log("Adding %d foreground seeds (distance: %f)", seeds.size(), distance);
		
		int size = (int) Math.min(distance/2, MAX_BRUSH_SIZE);
		
		AnnotationManager annotations = getAnnotations();
		for (Index2D seed : seeds) {
			this.seeds.add(seed);
			Point pt = getSwtPointForIndex(seed);
			annotations.add(new Annotation(AnnotationType.Foreground, size, pt));
		}
	}

	private static IntMatrix distanceTransform(ByteMatrix error) {
		DistanceTransform dt = new DistanceTransform();
		dt.init(error, Constants.ERROR_VALUE);
		return dt.computeSquareTransform();
	}
	
	public static boolean isNonDeterministic() {
		return false;
	}
}