package ie.dcu.apps.ist.exp;

import ie.dcu.eval.Evaluator;
import ie.dcu.segment.SegmentationMask;
import ie.dcu.segment.SegmentationContext;
import ie.dcu.apps.ist.AppWindow;
import ie.dcu.apps.ist.actions.*;

import java.io.*;
import java.util.List;

public class ExperimentManager {
	
	private final AppWindow window;
	
	private Experiment experiment;
	private SegmentationContext activeContext;
	private SegmentationMask referenceMask;
	private ExperimentResults results;
	private List<Task> tasks;
	private Task activeTask;
	private int progress;
	private boolean aborted;
	
	
	public ExperimentManager(AppWindow window) {	
		this.window = window;
	}
	
	
	public void beginExperiment(Experiment ex) throws IOException {
		experiment = ex;
		tasks = experiment.getTasks();
		activeContext = null;
		activeTask = null;
		progress = 0;
		results = new ExperimentResults(experiment.getOutputFile());
		results.beginDocument();
		results.beginExperiment(experiment);
	}
	
	
	public void endExperiment() throws IOException {
		results.endExperiment();
		results.flush();
		results.close();
	}
	
	
	public void abortExperiment() {
		results.endExperiment();
		results.close();
		aborted = true;
	}
	
	
	public void beginTask() throws IOException {
		if (isComplete()) {
			throw new IllegalStateException();
		}
		
		// Set active task
		activeTask = tasks.get(progress);
		
		// Load context into main window
		loadContext();

		// Load reference mask
		loadReferenceMask();
		
		// Write task header to results
		results.beginTask(progress, activeTask);
	}
	
	
	public void endTask() throws IOException {
		
		try {
			// Save final mask
			StorageSelection ss = experiment.getStorageSelection();
			if (ss != StorageSelection.None) {
				store(String.format("task-%d-final.png", progress));
			}
			
		} finally {
			
			activeTask = null;
			activeContext = null;
			progress++;
			results.endTask();
			results.flush();
		}
	}
	
	
	public boolean isComplete() {
		return tasks.size() == progress || aborted;
	}
	
	
	public Experiment getExperiment() {
		return experiment;
	}
	
	
	public String getTaskDescription() {
		return activeTask.getDescription();
	}
	
	
	public Task getActiveTask() {
		return activeTask;
	}
	
	
	public int getTaskCount() {
		return tasks.size();
	}
	
	
	public int getProgress() {
		return progress;
	}
	
	
	public SegmentationContext getActiveContext() {
		return activeContext;
	}
	
	
	public List<Evaluator> getEvaluators() {
		return experiment.getEvaluators();
	}
	
	
	public File getImageFile() {
		return activeTask.getImageFile();
	}
	
	
	public SegmentationMask getCurrentMask() {
		return activeContext.getMask();
	}
	
	
	public SegmentationMask getReferenceMask() {
		return referenceMask;
	}
	
	
	public void evaluate(int elapsed) {
		
		if (!getCurrentMask().isUnknown()) {
			SegmentationMask im = getCurrentMask();
			SegmentationMask gt = getReferenceMask();
			
			results.beginEvaluation(elapsed);
			
			for (Evaluator evaluator : experiment.getEvaluators()) {
				// Execute evaluator
				evaluator.run(im, gt);
				
				// Add results
				addResults(evaluator);
			}
			
			results.endEvaluation();
		}
	}
	
	
	public void store(int elapsed) throws IOException {
		// Save current mask
		if (!getCurrentMask().isUnknown()) {
			// Only save non empty masks
			StorageSelection ss = experiment.getStorageSelection();
			if (ss == StorageSelection.All) {
				store(String.format("task-%d-time-%d.png", progress, elapsed));
			}
		}
	}
	
	
	private void store(String name) throws IOException {
		SegmentationMask im = getCurrentMask();
		File dir = experiment.getStorageDirectory();
		File file = new File(dir, name);
		im.save(file);
	}


	private void addResults(Evaluator evaluator) {
		String[] measures = evaluator.getMeasures();
		for (String measure : measures) {
			double value = evaluator.getMeasure(measure);
			results.addMeasure(measure, value);
		}
	}


	private void loadReferenceMask() throws IOException {
		referenceMask = new SegmentationMask(activeContext.getBounds());
		referenceMask.load(activeTask.getMaskFile());
	}


	private void loadContext() throws IOException {
		ActionManager actions = window.getActions();
		OpenAction action = actions.get(OpenAction.class);
	
		if (!action.open(getImageFile())) {
			throw new IOException("Unable to load image file");
		}
		
		activeContext = window.getContext();
	}
}