DocAwareIterator.java

public class DocAwareIterator implements LabelAwareIterator{

	private AtomicInteger currPosition;
	
	private List<String> labels;
	private Map<String, String> docData;
	
	private LabelsSource generator;

	public DocAwareIterator(Map<String, String> docData) {
		this.currPosition = new AtomicInteger(0);
		this.generator = new LabelsSource();
		
		this.labels = new ArrayList<>(docData.keySet());
		this.docData = docData;
	}

	@Override
	public boolean hasNext() {
		 return hasNextDocument();
	}

	@Override
	public LabelledDocument next() {
		return nextDocument();
	}

	@Override
	public LabelsSource getLabelsSource() {
		return generator;
	}

	@Override
	public boolean hasNextDocument() {
		return currPosition.get() < labels.size();
	}

	@Override
	public LabelledDocument nextDocument() {
		LabelledDocument document = new LabelledDocument();

		
		String currentLabel = labels.get(currPosition.getAndIncrement());
		String currentContent = docData.get(currentLabel);
		
		document.setContent(currentContent);
		document.addLabel(currentLabel);
		
		generator.storeLabel(currentLabel);
		
		return document;
	}

	@Override
	public void reset() {
		currPosition.set(0);
		generator.reset();

	}

	@Override
	public void shutdown() { 
		// no-op 
	}



}