analysis/org.eclipse.stem.analysis/src/org/eclipse/stem/analysis/util/CSVScenarioStreamer.java

package org.eclipse.stem.analysis.util;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.StringTokenizer;

import org.eclipse.stem.analysis.LogInitializationException;

/*******************************************************************************
 * Copyright (c) 2007, 2008 IBM Corporation and others.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 *******************************************************************************/


/**
 * Adds methods for streaming data instead of loading all data files into memory
 */
public class CSVScenarioStreamer extends CSVscenarioLoader {

	public static final int NUM_LEAD_COLS = 2;
	
	// First Map: Key is population identifier (UNKNOWN_POP_IDENTIFIER_KEY if not known)
	// Second Map: Key is admin level
	// Third Map: Key is compartment
	
	Map<String, Map<Integer, Map<String, BufferedReader>>> streamReaders = new HashMap<String, Map<Integer, Map<String, BufferedReader>>>();
	Map<String, Map<Integer, Map<String, Integer>>> streamReaderCurrentRow = new HashMap<String, Map<Integer, Map<String, Integer>>>();
	Map<String, Map<Integer, Map<String, List<Double>>>> streamReadersLastRow = new HashMap<String, Map<Integer, Map<String, List<Double>>>>();
	
	Map<String, Map<Integer, List<String>>> headersMap = new HashMap<String, Map<Integer, List<String>>>();
	
	private int row = -1; // Which row
	
	/**
	 * @param path
	 */
	
	public CSVScenarioStreamer(String path) throws LogInitializationException {
		super(path);
	}
	
	
	/**
	 * prepareForStreaming
	 * 
	 * @param populationId Population identifier
	 * @param resolution The desired resolution, or -1 of all files should be streamed
	 */
	
	public void prepareForStreaming(String populationId, int resolution) throws IOException, LogInitializationException {
		
		if(!csvDataFiles.containsKey(populationId))
			throw new LogInitializationException(POPULATION_NOT_FOUND+populationId);
		
		List<File>fileList = csvDataFiles.get(populationId);
		
		for(int i = 0; i < fileList.size(); i ++) {
			File f = fileList.get(i);
			String name = f.getName();
			if(f.getName().startsWith(".")) continue;//$NON-NLS-1$
			int _idx = name.lastIndexOf('_');
			int dotidx = name.lastIndexOf('.');
			if(_idx <0 || dotidx < 0 || dotidx < _idx) continue; // not a data file
			String state = name.substring(0, _idx);

			// discovered a state
			registerKey(state);
			
			int res = Integer.parseInt(name.substring(_idx+1, dotidx));
			if(resolution != -1 && res != resolution) continue; // wrong resolution
			openFile(populationId, state, f, res);
		}
	}
	
	/**
	 * openFile. Open data file and prepare for streaming
	 * 
	 * @param population The population id
	 * @param state The state
	 * @param File file
	 * @param adminLevel
	 */
	
	public void openFile(String population, String state, File f, int adminLevel) throws IOException {
		BufferedReader reader = new BufferedReader(new FileReader(f));
		
		if(!streamReaders.containsKey(population)) {
			streamReaders.put(population, new HashMap<Integer, Map<String, BufferedReader>>());
			streamReaderCurrentRow.put(population, new HashMap<Integer, Map<String, Integer>>());
			streamReadersLastRow.put(population, new HashMap<Integer, Map<String, List<Double>>>());
		}
		if(!streamReaders.get(population).containsKey(adminLevel)) {
			streamReaders.get(population).put(adminLevel, new HashMap<String, BufferedReader>());
			streamReaderCurrentRow.get(population).put(adminLevel, new HashMap<String, Integer>());
			streamReadersLastRow.get(population).put(adminLevel, new HashMap<String, List<Double>>());
		}
		Map<String, BufferedReader>stateReaderMap = streamReaders.get(population).get(adminLevel);
		Map<String, Integer>stateRowMap = streamReaderCurrentRow.get(population).get(adminLevel);
		Map<String, List<Double>>stateRowData = streamReadersLastRow.get(population).get(adminLevel);
		stateReaderMap.put(state, reader);
		stateRowMap.put(state, -1); // no row read yet, first row is row 0
		stateRowData.put(state, null); // no data yet
		String header;
		
		if ( (header=reader.readLine()) != null ) { 
			if(!headersMap.containsKey(population))
				headersMap.put(population, new HashMap<Integer, List<String>>());
			Map<Integer, List<String>> adminMap = headersMap.get(population);
			if(!adminMap.containsKey(adminLevel)) {
				List<String>headerElements = new ArrayList<String>();
				adminMap.put(adminLevel, headerElements);
				StringTokenizer st = new StringTokenizer(header, ",");//$NON-NLS-1$
	
				int col = 0;
				while (st.hasMoreTokens()) {
					String val = st.nextToken().trim();
					if(col < NUM_LEAD_COLS) {++col;continue;}
					headerElements.add(val);
					++col;
				}
			}
		} // read the header
		/////////////////////
	}
	
	/**
	 * streamRow Stream next row of data
	 * 
	 * @param population The population id
	 * @param adminLevel which admin level
	 * @param state Which state (e.g. S)
	 * @return List<Double> The data, or null if end of line is encountered;
	 */
	
	public List<Double>streamRow(String population, int adminLevel, String state) throws IOException {
		BufferedReader br = streamReaders.get(population).get(adminLevel).get(state);
		if(br == null) return null; // stream is closed
		String data;
		if ( (data=br.readLine()) != null ) { 
			
			if(streamReadersLastRow.get(population).get(adminLevel).get(state) == null) 
				streamReadersLastRow.get(population).get(adminLevel).put(state, new ArrayList<Double>());
			List<Double>result = streamReadersLastRow.get(population).get(adminLevel).get(state);
			result.clear();
			StringTokenizer st = new StringTokenizer(data, ",");//$NON-NLS-1$
			
			int col = 0;
			while (st.hasMoreTokens()) {
				String val = st.nextToken().trim();
				if(col < NUM_LEAD_COLS) {++col;continue;}
				result.add(Double.parseDouble(val));
				++col;
			}
			streamReaderCurrentRow.get(population).get(adminLevel).put(state, streamReaderCurrentRow.get(population).get(adminLevel).get(state)+1);
			
			return result;
		}
		br.close();
		streamReaders.get(population).get(adminLevel).put(state, null);// gc reader
		return null; // end of file reached
	}
	
	public int getCurrentRow(String population, int adminLevel, String state) {
		return streamReaderCurrentRow.get(population).get(adminLevel).get(state);
	}
	
	public boolean containsState(String population, int adminLevel, String state) {
		return streamReaderCurrentRow.get(population).get(adminLevel).containsKey(state);
	}
	
	/**
	 * streamRow Stream next row of data
	 * 
	 * @param population Which population id
	 * @param adminLevel which admin level
	 * @param state Which state (e.g. S)
	 * @return List<Double> The data
	 */
	
	public List<Double>getLastFetchedRow(String population, int adminLevel, String state) {
		List<Double>result = streamReadersLastRow.get(population).get(adminLevel).get(state);
		return result;
	}

	/**
	 * Return the position of the specified key 
	 * @param population
	 * @param adminLevel
	 * @param key
	 * @return position
	 */
	public int getDataPosition(String population, int adminLevel, String key) {
		return headersMap.get(population).get(adminLevel).indexOf(key);
	}
	
	/**
	 * Close all open files
	 */
	public void close() throws IOException {
		if(streamReaders == null) return;
		for(Map.Entry<String, Map<Integer, Map<String, BufferedReader>>> entry:streamReaders.entrySet())
			for(Map.Entry<Integer, Map<String, BufferedReader>> entry2:entry.getValue().entrySet())
				for(Map.Entry<String, BufferedReader> entry3:entry2.getValue().entrySet())
					if(entry3.getValue() != null) entry3.getValue().close();
	}
	
}